fix(voice): preserve pre-decoder cancellation

Signed-off-by: John Tennant <jtennant@squareup.com>

crates/buzz-voice/src/pocket.rs

@@ -210,8 +210,9 @@ impl PocketTts {
     ///
     /// Callback sample buffers contain all PCM produced for this call so far.
     /// Their lengths never decrease, but equal lengths are allowed while the
-    /// engine advances between internal model-safe text chunks. Returning
-    /// `false` interrupts synthesis before the next decoder block.
+    /// engine advances before PCM is available or between internal model-safe
+    /// text chunks. Returning `false` interrupts synthesis before the next
+    /// model or decoder step.
     pub fn synth_chunk_streaming<F>(
         &self,
         text: &str,
@@ -346,6 +347,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn equal_length_pre_decoder_callback_can_cancel() {
+        let mut callback = |samples: &[f32], progress: f32| {
+            assert!(samples.is_empty());
+            assert_eq!(progress, 0.25);
+            false
+        };
+        let mut samples = Vec::new();
+
+        assert!(!append_and_callback(&mut samples, &[], &mut callback, 0.25)
+            .expect("pre-decoder cancellation callback"));
+    }
+
     #[test]
     fn callback_panic_is_reported_without_unwinding() {
         let mut callback = |_: &[f32], _: f32| -> bool {
@@ -432,7 +446,7 @@ mod tests {
 
     #[test]
     #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"]
-    fn production_streaming_callback_interrupts_after_first_decoder_block() {
+    fn production_streaming_callback_interrupts_before_pcm_is_available() {
         let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR")
             .expect("set BUZZ_POCKET_TEST_MODEL_DIR to an April INT8 model directory");
         let engine = load_text_to_speech(&dir).expect("load April INT8 engine");
@@ -450,7 +464,9 @@ mod tests {
                 "en",
                 &style,
                 1,
-                |_, _| {
+                |samples, progress| {
+                    assert!(samples.is_empty());
+                    assert!(progress <= 0.5);
                     callback_at = Some(started.elapsed());
                     false
                 },
@@ -465,4 +481,41 @@ mod tests {
 
         assert!(matches!(outcome, SynthesisOutcome::Interrupted));
     }
+
+    #[test]
+    #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"]
+    fn production_streaming_callback_interrupts_after_first_decoder_block() {
+        let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR")
+            .expect("set BUZZ_POCKET_TEST_MODEL_DIR to an April INT8 model directory");
+        let engine = load_text_to_speech(&dir).expect("load April INT8 engine");
+        let style = load_voice_style(&Path::new(&dir).join("reference_sample.wav"))
+            .expect("load reference voice");
+        let mut callback_at = None;
+        let started = std::time::Instant::now();
+
+        let outcome = engine
+            .synth_chunk_streaming(
+                "This sentence is long enough to require more than one decoder block.",
+                "en",
+                &style,
+                1,
+                |samples, progress| {
+                    if samples.is_empty() {
+                        return true;
+                    }
+                    assert!(progress > 0.5);
+                    callback_at = Some(started.elapsed());
+                    false
+                },
+            )
+            .expect("interrupt production streaming after decoded PCM");
+        let callback_at = callback_at.expect("decoder must produce a callback");
+        let cancellation_latency = started.elapsed().saturating_sub(callback_at);
+        eprintln!(
+            "decoded_callback_to_cancel_return_ms={:.1}",
+            cancellation_latency.as_secs_f64() * 1000.0
+        );
+
+        assert!(matches!(outcome, SynthesisOutcome::Interrupted));
+    }
 }

crates/buzz-voice/src/pocket_april.rs

@@ -285,7 +285,7 @@ impl AprilPocketTts {
         &mut self,
         prepared: &AprilPreparedPrompt,
         style: &VoiceStyle,
-        callback: F,
+        mut callback: F,
     ) -> Result<AprilSynthesisOutcome, String>
     where
         F: FnMut(&[f32], f32) -> bool,
@@ -316,9 +316,18 @@ impl AprilPocketTts {
         let text_embeddings = self.text_embeddings(token_ids)?;
         self.run_flow_main_prefix(&text_embeddings, &mut flow_state)?;
         let max_frames = estimate_max_frames(token_count, self.bundle.frame_rate);
-        let latents =
-            self.generate_latents(max_frames, prepared.frames_after_eos, &mut flow_state)?;
-        self.decode_latents(&latents, callback)
+        let (latents, interrupted) = self.generate_latents(
+            max_frames,
+            prepared.frames_after_eos,
+            &mut flow_state,
+            &mut callback,
+        )?;
+        if interrupted {
+            return Ok(AprilSynthesisOutcome::Interrupted);
+        }
+        self.decode_latents(&latents, |samples, progress| {
+            callback(samples, 0.5 + progress * 0.5)
+        })
     }
 
     fn prepared_token_count(&self, text: &str) -> Result<usize, String> {
@@ -478,18 +487,29 @@ impl AprilPocketTts {
         replace_state_from_outputs(state, &mut outputs)
     }
 
-    fn generate_latents(
+    fn generate_latents<F>(
         &mut self,
         max_frames: usize,
         frames_after_eos: usize,
         state: &mut [StateValue],
-    ) -> Result<Vec<f32>, String> {
+        callback: &mut F,
+    ) -> Result<(Vec<f32>, bool), String>
+    where
+        F: FnMut(&[f32], f32) -> bool,
+    {
         let mut current = vec![f32::NAN; self.bundle.latent_dim];
         let mut latents = Vec::with_capacity(max_frames * self.bundle.latent_dim);
         let mut eos_step = None;
         let mut rng = rand::rng();
 
         for step in 0..max_frames {
+            // Preserve the mobile callback's pre-PCM cancellation point while
+            // reserving the second half of progress for decoder-block output.
+            // The empty block becomes an equal-length cumulative callback in
+            // `PocketTts::synth_chunk_streaming`.
+            if !callback(&[], step as f32 / max_frames as f32 * 0.5) {
+                return Ok((latents, true));
+            }
             let sequence = Tensor::from_array((
                 vec![1_i64, 1, self.bundle.latent_dim as i64],
                 current.clone().into_boxed_slice(),
@@ -575,7 +595,7 @@ impl AprilPocketTts {
             current.clone_from(&noise);
             latents.extend_from_slice(&noise);
         }
-        Ok(latents)
+        Ok((latents, false))
     }
 
     fn decode_latents<F>(