jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-28 13:01:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(huddle): Pocket TTS quality overhaul — reference parity + cross-message pipelining (#997)

Browse Source 
Signed-off-by: Tyler Longwell <tlongwell@squareup.com>
Co-authored-by: npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d <011987e296fd5006292d2f930b574be47c7801048d1983c46c425d3c95f0cffd@sprout-oss.stage.blox.sqprod.co>
This commit is contained in:
tlongwell-block
2026-06-12 12:56:42 -04:00
committed by GitHub
parent 00dc4915d4
commit 12433077ad
8 changed files with 1367 additions and 1419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
desktop/scripts/check-file-sizes.mjs
View File

@@ -35,7 +35,6 @@ const overrides = new Map([
["src-tauri/src/managed_agents/runtime.rs", 1940],
["src-tauri/src/managed_agents/personas.rs", 1080],
["src-tauri/src/managed_agents/persona_card.rs", 1050],
["src-tauri/src/huddle/tts.rs", 1364],
["src/shared/api/tauri.ts", 1196],
["src-tauri/src/nostr_convert.rs", 1126],
["src/shared/api/relayClientSession.ts", 1022],

2
desktop/src-tauri/examples/pocket_bench.rs
View File

@@ -63,7 +63,7 @@ fn main() {
let gen = || GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(samples.clone()),
reference_sample_rate: sr,
..Default::default()

123
desktop/src-tauri/examples/pocket_clip_probe.rs Normal file
View File

@@ -0,0 +1,123 @@
//! Clipping probe for any fixed playback gain applied after Pocket TTS synth.
//!
//! Synthesises a spread of sentences (short/long, calm/energetic) and reports
//! the raw peak of each, the post-gain peak, and the fraction of samples that
//! would hit a ±1.0 clamp — i.e. how much a fixed gain would flat-top the
//! waveform ("blown out" distortion).
//!
//! History: the production pipeline briefly shipped a fixed 9.3× gain
//! calibrated on a single bench utterance that peaked at 0.076. This probe
//! showed real output peaks at 0.4–0.97, so that gain clipped 13–34% of all
//! samples (the 2026-06-12 "blown out" report). Production now applies no
//! gain — run this probe before reintroducing one.
//!
//! Run with model files in ~/.buzz/models/pocket-tts (override with arg 1):
//! cargo run --release --example pocket_clip_probe
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
/// Candidate gain under test (the regressed production value).
const GAIN: f32 = 9.3;
const PROMPTS: &[&str] = &[
"Hello, this is a test of the new Pocket TTS engine running on sherpa-onnx.",
"Yep, I can hear you.",
"Absolutely! That sounds fantastic, let's do it right now!",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I found three problems in the code: a race condition, a memory leak, and an off-by-one error in the loop bounds.",
"No.",
"Warning! The build failed because seventeen tests crashed unexpectedly!",
"Sure, I can walk you through the whole pipeline step by step whenever you're ready.",
];
fn main() {
let model_dir = std::env::args().nth(1).unwrap_or_else(|| {
dirs::home_dir()
.expect("home dir")
.join(".buzz/models/pocket-tts")
.to_string_lossy()
.into_owned()
});
eprintln!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let cfg = OfflineTtsConfig {
model: OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
},
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let voice_samples = wave.samples().to_vec();
let voice_sr = wave.sample_rate();
let gen = || GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 1.0,
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
// Warmup.
let _ = engine.generate_with_config("warmup.", &gen(), None::<fn(&[f32], f32) -> bool>);
println!(
"{:<46} | {:>8} | {:>9} | {:>9} | {:>10}",
"prompt", "raw peak", "raw RMS", "post-gain", "% clipped"
);
println!("{}", "-".repeat(95));
let mut worst_clip = 0.0f32;
for prompt in PROMPTS {
let out = engine
.generate_with_config(prompt, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("synth");
let samples = out.samples();
let peak = samples.iter().fold(0.0f32, |m, s| m.max(s.abs()));
let rms = (samples.iter().map(|s| s * s).sum::<f32>() / samples.len() as f32).sqrt();
let post = peak * GAIN;
let clipped = samples.iter().filter(|s| s.abs() * GAIN > 1.0).count();
let clip_pct = 100.0 * clipped as f32 / samples.len() as f32;
worst_clip = worst_clip.max(clip_pct);
let label: String = prompt.chars().take(44).collect();
println!("{label:<46} | {peak:>8.4} | {rms:>9.4} | {post:>9.3} | {clip_pct:>9.3}%");
}
println!();
println!(
"Verdict: worst-case clipped fraction {worst_clip:.3}% — {}",
if worst_clip > 0.1 {
"AUDIBLE DISTORTION LIKELY (gain too hot)"
} else if worst_clip > 0.0 {
"marginal — occasional transient clipping"
} else {
"no clipping at this gain"
}
);
}

4
desktop/src-tauri/examples/pocket_onset_probe.rs
View File

@@ -72,7 +72,7 @@ fn main() {
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
@@ -100,7 +100,7 @@ fn main() {
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()

220
desktop/src-tauri/examples/prod_probe.rs
View File

@@ -1,220 +0,0 @@
//! Reproduction probe for the Pocket TTS first-phoneme-drop fix.
//!
//! Pocket TTS' FlowLM has an autoregressive cold-start that occasionally
//! smears or drops the first phoneme of short utterances — see
//! kyutai-labs/pocket-tts #91, #70 and sherpa-onnx #3180. The cure is to
//! prepend a sacrificial `". . "` cold-start absorber to short prompts
//! and trim the resulting leading audio. This example reproduces both
//! variants of generated audio so you can listen-test the fix at the
//! exact `GenerationConfig` we ship in production (`huddle::pocket`):
//!
//! - silence_scale: 0.0 (production)
//! - max_frames: 100 (short) / sherpa default 500 (long)
//! - num_steps: 1
//! - speed: 1.05
//!
//! Note: production does NOT override `frames_after_eos` — sherpa-onnx's
//! default of 3 is what we want. The previous attempt to override it for
//! long prompts caused the "first 'yep' is static" regression (commit
//! 1dbfa2c). This probe mirrors that decision.
//!
//! Run:
//! cargo run --release --example prod_probe
//! cargo run --release --example prod_probe /path/to/pocket-tts
//!
//! Output (per (label, seed) pair):
//! /tmp/prod_<label>_s<seed>_raw.wav — raw engine output
//! /tmp/prod_<label>_s<seed>_trimmed.wav — post-trim (what production
//! ships, for `_sac` labels)
//!
//! The "no sacrificial" variants have no trim applied (just _raw); they
//! show what production produces for long prompts. The "_sac" variants
//! show both raw and trimmed, which is what `huddle::pocket::synth_chunk`
//! returns for short prompts. Listen back with `afplay`.
use std::collections::HashMap;
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
const SHORT_PROMPT_MAX_FRAMES: i32 = 100;
// Mirror of huddle::pocket trim constants so the probe stays in sync with
// production. If you change either side, change both.
const TRIM_SCAN_START_SAMPLES: usize = (SAMPLE_RATE as usize * 30) / 1000;
const TRIM_SILENCE_THRESHOLD: f32 = 0.02;
const TRIM_MIN_GAP_SAMPLES: usize = (SAMPLE_RATE as usize * 50) / 1000;
const TRIM_MAX_DROP_SAMPLES: usize = (SAMPLE_RATE as usize * 1200) / 1000;
/// Mirror of `huddle::pocket::trim_leading_cold_start` — keep in sync.
fn trim_leading_cold_start(samples: &mut Vec<f32>) {
if samples.len() <= TRIM_SCAN_START_SAMPLES {
return;
}
let mut silence_run_start: Option<usize> = None;
let mut gap_end: Option<usize> = None;
for (i, sample) in samples.iter().enumerate().skip(TRIM_SCAN_START_SAMPLES) {
if sample.abs() < TRIM_SILENCE_THRESHOLD {
silence_run_start.get_or_insert(i);
} else if let Some(start) = silence_run_start {
if i - start >= TRIM_MIN_GAP_SAMPLES {
gap_end = Some(i);
break;
}
silence_run_start = None;
}
}
let Some(end) = gap_end else {
return;
};
if end > TRIM_MAX_DROP_SAMPLES {
return;
}
samples.drain(..end);
}
// (label, raw_text_before_prep, sacrificial_prefix_to_add_after_pad)
const TESTS: &[(&str, &str, &str)] = &[
// Short, previously-failing — sacrificial applied
("imhappy_sac", "I'm happy.", ". . "),
("imsorry_sac", "I'm sorry.", ". . "),
("imready_sac", "I'm ready.", ". . "),
// Short, previously-OK — sacrificial applied
("yep_sac", "Yep.", ". . "),
("isee_sac", "I see you.", ". . "),
// Long — sacrificial NOT applied (per design)
("longer_nosac", "Hello, how can I help you today?", ""),
("multi_nosac", "Yes, that works. Let me try again.", ""),
];
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let cfg = OfflineTtsConfig {
model: OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
},
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let ref_samples = wave.samples().to_vec();
let ref_sr = wave.sample_rate();
let seeds: &[i32] = &[42, 1337, 99999, 7, 314159];
println!(
"{:18} | {:>6} | {:>5} | {:>7} | gap_search (50ms,0.02) | path",
"test", "seed", "kind", "len_ms"
);
println!("{}", "-".repeat(110));
for (label, raw_text, sacrificial) in TESTS {
// Mirror huddle::pocket::prepare_pocket_prompt:
// - cleaned text starts with capital, ends with punctuation
// - short (≤4 words) → pad + sacrificial; max_frames=100
// - long → unchanged; no max_frames override
let cleaned = raw_text.trim();
let word_count = cleaned.split_whitespace().count();
let is_short = word_count <= 4;
let pad = if is_short { " " } else { "" };
let prompt = format!("{pad}{sacrificial}{cleaned}");
// Only short prompts get the post-synth trim in production; long
// prompts pass through unmodified.
let trim_in_production = is_short && !sacrificial.is_empty();
for seed in seeds {
let mut extra: HashMap<String, serde_json::Value> = HashMap::new();
extra.insert("seed".to_string(), serde_json::Value::from(*seed));
if is_short {
extra.insert(
"max_frames".to_string(),
serde_json::Value::from(SHORT_PROMPT_MAX_FRAMES),
);
}
let gen = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 0.0, // PRODUCTION SETTING
reference_audio: Some(ref_samples.clone()),
reference_sample_rate: ref_sr,
extra: Some(extra),
..Default::default()
};
let audio = engine
.generate_with_config(&prompt, &gen, None::<fn(&[f32], f32) -> bool>)
.expect("synth");
// Raw output (what the engine returned).
let raw_samples = audio.samples().to_vec();
let raw_ms = raw_samples.len() as f32 / SAMPLE_RATE as f32 * 1000.0;
let raw_gap = find_gap(&raw_samples, SAMPLE_RATE, 0.02, 50);
let raw_path = format!("/tmp/prod_{}_s{}_raw.wav", label, seed);
sherpa_onnx::write(&raw_path, &raw_samples, SAMPLE_RATE as i32);
println!(
"{:18} | {:>6} | {:>5} | {:>5.0}ms | {:>22} | {}",
label, seed, "raw", raw_ms, raw_gap, raw_path
);
// Trimmed output — what synth_chunk actually returns to tts.rs
// for short prompts. For long prompts the engine output is
// returned untrimmed, so we skip writing a separate file.
if trim_in_production {
let mut trimmed = raw_samples;
trim_leading_cold_start(&mut trimmed);
let trimmed_ms = trimmed.len() as f32 / SAMPLE_RATE as f32 * 1000.0;
let trimmed_path = format!("/tmp/prod_{}_s{}_trimmed.wav", label, seed);
sherpa_onnx::write(&trimmed_path, &trimmed, SAMPLE_RATE as i32);
println!(
"{:18} | {:>6} | {:>5} | {:>5.0}ms | {:>22} | {}",
label, seed, "trim", trimmed_ms, "(post-trim)", trimmed_path
);
}
}
println!();
}
}
fn find_gap(samples: &[f32], sr: u32, thresh: f32, min_ms: u32) -> String {
let scan_start = (sr as usize * 30) / 1000;
let min_samples = (sr as usize * min_ms as usize) / 1000;
let mut silence_from: Option<usize> = None;
for (i, sample) in samples.iter().enumerate().skip(scan_start) {
if sample.abs() < thresh {
silence_from.get_or_insert(i);
} else if let Some(start) = silence_from {
if i - start >= min_samples {
return format!(
"{:.0}..{:.0}ms ({:.0}ms)",
start as f32 / sr as f32 * 1000.0,
i as f32 / sr as f32 * 1000.0,
(i - start) as f32 / sr as f32 * 1000.0
);
}
silence_from = None;
}
}
"<no gap>".to_string()
}

345
desktop/src-tauri/src/huddle/pocket.rs
View File

@@ -81,10 +81,17 @@ const VOICE_EMBEDDING_CACHE_CAPACITY: i32 = 16;
/// synthesis time by ~5× with no audible benefit on this model.
const SYNTH_NUM_STEPS: i32 = 1;
/// Disable the upstream default 200 ms of pre/post silence padding. We splice
/// `INTER_SENTENCE_SILENCE` in `tts.rs` ourselves and don't want a double
/// helping of leading silence on every utterance.
const SYNTH_SILENCE_SCALE: f32 = 0.0;
/// Leave the generated audio's silences untouched (1.0 is the identity).
///
/// sherpa-onnx's `ScaleSilence` (`offline-tts.cc`) is *not* pre/post padding
/// control: it finds every interior silence run ≥ 0.2 s (|s| ≤ 0.01) and
/// multiplies its length by this factor. The previous value of 0.0 — set
/// under the mistaken belief it disabled lead-in/lead-out padding — deleted
/// every natural pause inside an utterance: clause breaks, breaths, the gap
/// after a comma. Words slammed together and endings cut abruptly. The
/// reference Pocket TTS pipeline does not post-process silence at all;
/// 1.0 restores parity.
const SYNTH_SILENCE_SCALE: f32 = 1.0;
/// sherpa-onnx upstream default for `max_frames` (LM steps), in
/// `offline-tts-pocket-impl.h:Generate`. 500 steps ≈ 40 s of audio at the
@@ -112,61 +119,19 @@ const SHORT_PROMPT_WORD_THRESHOLD: usize = 4;
/// Number of leading spaces prepended to short prompts. The upstream Python
/// uses exactly 8 — keep parity rather than tuning blindly.
///
/// This is upstream's *only* mitigation for the FlowLM cold-start smear on
/// short utterances (kyutai-labs/pocket-tts #91, #70): the autoregressive
/// generation has a 2–3 step "settle" period where the first phoneme can be
/// smeared. A previous revision added a sacrificial `". . "` prefix plus an
/// amplitude-threshold trim to strip the rendered prefix from the output —
/// but the trim's absolute threshold (0.02 against raw peaks of ~0.076) sat
/// in soft-onset territory and could eat real word starts, and its tuning
/// was calibrated against `silence_scale = 0.0` audio. Deleted in favour of
/// upstream parity: accept the occasional smeared first syllable rather
/// than risk trimming real speech.
const SHORT_PROMPT_PAD_SPACES: usize = 8;
/// Sacrificial cold-start prefix appended *after* the leading space pad for
/// short prompts. Pocket TTS' FlowLM autoregressive generation has a 2–3
/// step "settle" period at the start where the first generated phoneme can
/// be smeared or dropped entirely (see kyutai-labs/pocket-tts #91, #70 and
/// sherpa-onnx #3180). For short utterances like "I'm happy." the first
/// phoneme is most of the first word — losing it produces "m happy".
///
/// Two periods separated by a space act as a "phantom utterance" that the
/// model commits to, absorbing the cold-start. The pair (rather than a
/// single period) was empirically the only variant in our probe — see
/// `examples/prod_probe.rs` — that produced a usable post-sacrificial
/// silence gap on every random seed. The resulting leading sacrificial
/// audio is then stripped from the output by [`trim_leading_cold_start`]
/// before the buffer is returned to the synth pipeline.
///
/// Long prompts (>4 words) don't need this — the first phoneme already has
/// enough downstream context to avoid the smear, and an early natural pause
/// (e.g. the comma in "Hello, how can I help you?") could be misdetected as
/// the trim boundary.
const SACRIFICIAL_PREFIX: &str = ". . ";
// ── Leading cold-start trim (post-synth) ──────────────────────────────────────
/// Skip this many samples at the start of the synth buffer before looking
/// for the sacrificial→main silence gap. The Mimi decoder cold-start
/// produces ~30 ms of low-amplitude noise that we *don't* want to treat as
/// the gap. 30 ms × 24 kHz = 720 samples.
const TRIM_SCAN_START_SAMPLES: usize = (SAMPLE_RATE as usize * 30) / 1000;
/// Amplitude threshold below which a sample is considered "silence" for the
/// purposes of finding the post-sacrificial gap. Tuned empirically against
/// production-config probe data — the engine's own `ScaleSilence` uses 0.01,
/// but our boundary detection wants a looser threshold so that the breath /
/// aspiration of the rendered periods (which sits around 0.005–0.015) is
/// treated as silence too.
const TRIM_SILENCE_THRESHOLD: f32 = 0.02;
/// A silence run must be at least this many samples long to be accepted as
/// the sacrificial→main word boundary. 50 ms is comfortably longer than
/// inter-syllable silence within a normal word at this speed (typically
/// 10–30 ms), so this guards against trimming into the middle of the real
/// utterance. 50 ms × 24 kHz = 1200 samples.
const TRIM_MIN_GAP_SAMPLES: usize = (SAMPLE_RATE as usize * 50) / 1000;
/// Hard cap on how much audio we'll trim from the start. Production probe
/// data (with `silence_scale = 0.0`) shows valid trim boundaries land
/// between 30 ms and ~450 ms; 1.2 s is a wide safety margin. If the
/// detector finds a "gap" past this point it's almost certainly an interior
/// pause inside an unusually long short-prompt utterance — bail out and
/// emit untrimmed audio rather than corrupt it. 1.2 s × 24 kHz = 28800
/// samples.
const TRIM_MAX_DROP_SAMPLES: usize = (SAMPLE_RATE as usize * 1200) / 1000;
/// sherpa-onnx's documented `frames_after_eos` default. We deliberately do
/// *not* override this knob — the previous attempt to bump it for short
/// inputs and lower it for long inputs lowered it below the upstream default
@@ -290,20 +255,13 @@ pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct PreparedPrompt {
/// Text to hand to `OfflineTts::generate_with_config`. Capitalized,
/// punctuation-terminated, and (for short inputs) left-padded with spaces
/// plus a sacrificial `". . "` cold-start prefix.
/// punctuation-terminated, and (for short inputs) left-padded with
/// spaces — upstream's mitigation for the FlowLM cold-start smear.
pub text: String,
/// Value to pass via `GenerationConfig.extra["max_frames"]`, or `None` to
/// keep the upstream default of 500 LM steps. We only override on short
/// padded prompts where we have a tight expectation on output length.
pub max_frames: Option<i32>,
/// `true` iff this prompt received the short-input treatment (leading
/// space pad + sacrificial `". . "` prefix). The synth pipeline uses
/// this to decide whether to apply [`trim_leading_cold_start`] to the
/// output: long prompts have no sacrificial audio to strip, and
/// trimming them risks deleting real speech at a natural early pause
/// (e.g. the comma in "Hello, how can I help you?").
pub is_short: bool,
}
/// Mirror of the *text-preparation* half of upstream
@@ -318,10 +276,9 @@ pub(crate) struct PreparedPrompt {
/// 2. Capitalize the first letter.
/// 3. Append `.` if the text doesn't end in punctuation.
/// 4. If fewer than five words, prepend `SHORT_PROMPT_PAD_SPACES` spaces
/// followed by [`SACRIFICIAL_PREFIX`] (a `". . "` cold-start absorber —
/// see its docstring for the bug it works around), and return a tight
/// [`SHORT_PROMPT_MAX_FRAMES`] cap so the LM can't run away if EOS
/// still doesn't fire.
/// (upstream's cold-start mitigation — see the constant's docstring) and
/// return a tight [`SHORT_PROMPT_MAX_FRAMES`] cap so the LM can't run
/// away if EOS still doesn't fire.
///
/// We do **not** override `frames_after_eos` — sherpa-onnx's default of 3
/// is what we want. An earlier version set it to 1 on long inputs, which
@@ -377,16 +334,12 @@ pub(crate) fn prepare_pocket_prompt(input: &str) -> Option<PreparedPrompt> {
// Word count of the *cleaned but not padded* text — padding is whitespace
// only and would just lie to the threshold check below.
let word_count = cleaned.split_whitespace().count();
let is_short = word_count <= SHORT_PROMPT_WORD_THRESHOLD;
let (final_text, max_frames) = if is_short {
let mut padded = String::with_capacity(
cleaned.len() + SHORT_PROMPT_PAD_SPACES + SACRIFICIAL_PREFIX.len(),
);
let (final_text, max_frames) = if word_count <= SHORT_PROMPT_WORD_THRESHOLD {
let mut padded = String::with_capacity(cleaned.len() + SHORT_PROMPT_PAD_SPACES);
for _ in 0..SHORT_PROMPT_PAD_SPACES {
padded.push(' ');
}
padded.push_str(SACRIFICIAL_PREFIX);
padded.push_str(&cleaned);
(padded, Some(SHORT_PROMPT_MAX_FRAMES))
} else {
@@ -400,70 +353,9 @@ pub(crate) fn prepare_pocket_prompt(input: &str) -> Option<PreparedPrompt> {
Some(PreparedPrompt {
text: final_text,
max_frames,
is_short,
})
}
/// Strip the leading "sacrificial" audio produced by the `". . "` cold-start
/// prefix from a short-prompt synthesis result. Only call this when
/// [`PreparedPrompt::is_short`] is `true` — the trim looks for a long
/// silence run at the head of the buffer, and an early natural pause inside
/// a long unsacrificed utterance (e.g. the comma in "Hello, how can I help
/// you?") would be misclassified as the sacrificial gap.
///
/// Algorithm:
/// 1. Skip the first [`TRIM_SCAN_START_SAMPLES`] of the buffer (Mimi
/// cold-start noise we shouldn't classify as silence).
/// 2. Scan forward for the first run of samples below
/// [`TRIM_SILENCE_THRESHOLD`] that lasts at least
/// [`TRIM_MIN_GAP_SAMPLES`] — that's the post-sacrificial boundary.
/// 3. If that boundary lies beyond [`TRIM_MAX_DROP_SAMPLES`], treat it as
/// "almost certainly an interior pause" and *do not trim* — the safe
/// fallback is to play the slightly-degraded raw audio rather than
/// delete real speech.
/// 4. Otherwise, drop the leading samples up to the end of the silence
/// run. We don't insert a zero lead-in here — `tts.rs` owns playback
/// cushioning by prepending `SENTENCE_LEAD_IN_SAMPLES` of zeros before
/// each appended sentence chunk.
///
/// If the scan never finds a long-enough gap (≈1% of generations in the
/// production-config probe), the function is a no-op — the model trajectory
/// missed the expected sacrificial→main structure and we'd rather play the
/// raw buffer than emit silence.
fn trim_leading_cold_start(samples: &mut Vec<f32>) {
if samples.len() <= TRIM_SCAN_START_SAMPLES {
return;
}
let mut silence_run_start: Option<usize> = None;
let mut gap_end: Option<usize> = None;
for (i, sample) in samples.iter().enumerate().skip(TRIM_SCAN_START_SAMPLES) {
if sample.abs() < TRIM_SILENCE_THRESHOLD {
silence_run_start.get_or_insert(i);
} else if let Some(start) = silence_run_start {
if i - start >= TRIM_MIN_GAP_SAMPLES {
gap_end = Some(i);
break;
}
silence_run_start = None;
}
}
let Some(end) = gap_end else {
// No gap found — model didn't produce the expected sacrificial→main
// structure. Bail out and let the caller play the raw buffer.
return;
};
if end > TRIM_MAX_DROP_SAMPLES {
// Boundary too far into the audio to plausibly be the sacrificial
// gap. Almost certainly an interior pause in the real utterance —
// leave the audio alone.
return;
}
samples.drain(..end);
}
/// Build the `GenerationConfig.extra` HashMap from a [`PreparedPrompt`].
///
/// Centralised so the regression test below can assert that we **never**
@@ -544,16 +436,7 @@ impl PocketTts {
);
}
let mut samples = audio.samples().to_vec();
// For short prompts the prepared text includes a sacrificial ". . "
// prefix to absorb FlowLM/Mimi cold-start (see `SACRIFICIAL_PREFIX`).
// Strip the leading sacrificial audio before returning. Long prompts
// are never trimmed — they have no sacrificial audio, and an early
// natural pause could be mis-detected as the trim boundary.
if prepared.is_short {
trim_leading_cold_start(&mut samples);
}
Ok(samples)
Ok(audio.samples().to_vec())
}
}
@@ -571,23 +454,20 @@ mod tests {
}
/// Helper: the exact leading sequence prepended to every short prompt —
/// 8 spaces of padding followed by the sacrificial `". . "` cold-start
/// absorber. Centralising this keeps the assertions readable.
/// 8 spaces of padding (upstream's cold-start mitigation).
/// Centralising this keeps the assertions readable.
fn short_prefix() -> String {
let mut s = " ".repeat(SHORT_PROMPT_PAD_SPACES);
s.push_str(SACRIFICIAL_PREFIX);
s
" ".repeat(SHORT_PROMPT_PAD_SPACES)
}
#[test]
fn prepare_prompt_pads_and_capitalizes_one_word() {
// The "yep" case Tyler hit in production — bare lowercase one-word
// utterance with no punctuation. Must be padded with the short-prompt
// prefix (8 spaces + ". . " sacrificial), capitalized, terminated,
// with a tight `max_frames` cap to bound runaway gen.
// space pad, capitalized, terminated, with a tight `max_frames` cap
// to bound runaway gen.
let out = prepare_pocket_prompt("yep").expect("non-empty");
assert_eq!(out.text, format!("{}Yep.", short_prefix()));
assert!(out.is_short, "1-word input is short");
assert_eq!(out.max_frames, Some(SHORT_PROMPT_MAX_FRAMES));
const {
assert!(
@@ -607,31 +487,21 @@ mod tests {
#[test]
fn prepare_prompt_threshold_is_inclusive_at_four_words() {
// 4 words = short (padded + sacrificial + tight max_frames); 5 words
// = long (no padding, no sacrificial, no overrides — upstream
// defaults stand).
// 4 words = short (padded + tight max_frames); 5 words = long
// (no padding, no overrides — upstream defaults stand).
let four = prepare_pocket_prompt("one two three four").expect("non-empty");
assert!(four.is_short, "four-word input should be short");
assert!(
four.text.starts_with(' '),
"four-word input should start with the space pad"
);
assert!(
four.text.contains(SACRIFICIAL_PREFIX),
"four-word input should contain the sacrificial prefix"
assert_eq!(
four.text,
format!("{}One two three four.", short_prefix()),
"four-word input should get exactly the space pad"
);
assert_eq!(four.max_frames, Some(SHORT_PROMPT_MAX_FRAMES));
let five = prepare_pocket_prompt("one two three four five").expect("non-empty");
assert!(!five.is_short, "five-word input should NOT be short");
assert!(
!five.text.starts_with(' '),
"five-word input should NOT be padded"
);
assert!(
!five.text.contains(SACRIFICIAL_PREFIX),
"five-word input must not receive the sacrificial prefix"
);
assert_eq!(
five.max_frames, None,
"long inputs must leave sherpa-onnx's max_frames default in place"
@@ -642,9 +512,7 @@ mod tests {
fn prepare_prompt_does_not_pad_long_text() {
let long = "This is a longer sentence that the model should handle just fine.";
let out = prepare_pocket_prompt(long).expect("non-empty");
assert!(!out.is_short);
assert!(!out.text.starts_with(' '));
assert!(!out.text.contains(SACRIFICIAL_PREFIX));
assert_eq!(out.max_frames, None);
assert!(out.text.ends_with('.'));
}
@@ -652,8 +520,7 @@ mod tests {
#[test]
fn prepare_prompt_collapses_whitespace() {
let out = prepare_pocket_prompt("Hello world\n\nfriend").expect("non-empty");
// 3 words → short → padded + sacrificial. Interior whitespace
// collapsed.
// 3 words → short → padded. Interior whitespace collapsed.
assert_eq!(out.text, format!("{}Hello world friend.", short_prefix()));
}
@@ -671,16 +538,22 @@ mod tests {
assert!(out.text.contains("Дa."));
}
/// REGRESSION GUARD: short prompts must receive *only* whitespace
/// padding — no sacrificial text. A previous revision prepended a
/// `". . "` cold-start absorber and trimmed the rendered audio back out
/// with an amplitude threshold that could eat soft word onsets. If
/// non-whitespace ever reappears in the pad, the synth output will
/// contain audio for text the user never wrote.
#[test]
fn prepare_prompt_inserts_sacrificial_prefix_only_for_short() {
// Pinning the exact ordering: pad, then ". . ", then cleaned text.
// If this ever flips, the trim algorithm's calibration breaks.
fn prepare_prompt_pad_is_whitespace_only() {
let out = prepare_pocket_prompt("I'm happy.").expect("non-empty");
assert!(out.is_short);
let pad = " ".repeat(SHORT_PROMPT_PAD_SPACES);
let expected = format!("{pad}. . I'm happy.");
assert_eq!(out.text, expected);
assert_eq!(SACRIFICIAL_PREFIX, ". . ");
let pad_len = out.text.len() - "I'm happy.".len();
assert!(
out.text[..pad_len].chars().all(|c| c == ' '),
"short-prompt pad must be spaces only, got {:?}",
&out.text[..pad_len]
);
assert_eq!(out.text, format!("{}I'm happy.", short_prefix()));
}
// ── build_generation_extra ───────────────────────────────────────────────
@@ -769,110 +642,4 @@ mod tests {
assert!(SHORT_PROMPT_MAX_FRAMES >= 50, "would risk truncation");
}
}
// ── trim_leading_cold_start ──────────────────────────────────────────────
/// Build a synthetic buffer shaped like a sacrificial-prefixed Pocket TTS
/// output: a bit of "sacrificial" energy at the head (modelled as
/// alternating ±0.05 — above [`TRIM_SILENCE_THRESHOLD`] so it isn't
/// classified as silence, matching what real probe WAVs look like in the
/// 0–50 ms window), then a flat silence of `gap_ms` ms, then `tail_ms`
/// of "real speech" at peak `tail_peak`.
fn synth_buffer(sacrificial_ms: u32, gap_ms: u32, tail_ms: u32, tail_peak: f32) -> Vec<f32> {
let sr = SAMPLE_RATE as usize;
let mut v = Vec::new();
for i in 0..(sr * sacrificial_ms as usize / 1000) {
v.push(if i % 2 == 0 { 0.05 } else { -0.05 });
}
// Silence gap (true zeros — below TRIM_SILENCE_THRESHOLD).
v.extend(std::iter::repeat_n(0.0_f32, sr * gap_ms as usize / 1000));
// Real speech.
for i in 0..(sr * tail_ms as usize / 1000) {
v.push(if i % 2 == 0 { tail_peak } else { -tail_peak });
}
v
}
#[test]
fn trim_strips_sacrificial_and_keeps_only_speech() {
// 60 ms sacrificial + 100 ms gap + 500 ms speech at peak 0.3.
// After trim, the output is just the speech tail.
let mut v = synth_buffer(60, 100, 500, 0.3);
trim_leading_cold_start(&mut v);
// First sample should be speech (|s| ≥ 0.2). No zero lead-in here
// because tts.rs's `first_append` lead-in handles the device cushion.
assert!(
v[0].abs() > 0.2,
"first sample after trim should be speech, got {}",
v[0]
);
let actual_ms = (v.len() as f32 / SAMPLE_RATE as f32) * 1000.0;
assert!(
(actual_ms - 500.0).abs() < 5.0,
"expected ~500 ms of trimmed audio, got {actual_ms} ms"
);
}
#[test]
fn trim_is_noop_when_no_long_silence_gap_exists() {
// Pure speech: every sample is real (no gap >= 50 ms). Trimmer must
// leave the buffer untouched so we don't truncate the utterance.
let mut v = synth_buffer(0, 0, 600, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "no gap → no trim");
}
#[test]
fn trim_is_noop_when_gap_is_shorter_than_threshold() {
// 40 ms gap is below TRIM_MIN_GAP_SAMPLES (50 ms). Must not trigger.
let mut v = synth_buffer(60, 40, 600, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "sub-threshold gap → no trim");
}
#[test]
fn trim_is_noop_when_gap_is_beyond_max_drop_bound() {
// Gap starts at 1500 ms (past TRIM_MAX_DROP_SAMPLES = 1200 ms).
// This represents an interior pause inside an unusually long
// utterance that slipped past the short-prompt predicate; we must
// not chop the first 1.5 s of real audio.
let mut v = synth_buffer(1500, 200, 400, 0.3);
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before, "gap past max-drop bound → no trim");
}
#[test]
fn trim_is_noop_on_buffer_smaller_than_scan_start() {
// 20 ms buffer is smaller than TRIM_SCAN_START_SAMPLES (30 ms).
// Trimmer must early-return without panicking.
let mut v = vec![0.5f32; (SAMPLE_RATE as usize * 20) / 1000];
let before = v.clone();
trim_leading_cold_start(&mut v);
assert_eq!(v, before);
}
#[test]
fn trim_constants_use_sane_units() {
// Pin the constants in milliseconds so anyone tuning later can see
// at a glance what they're changing.
assert_eq!(
TRIM_SCAN_START_SAMPLES,
(SAMPLE_RATE as usize * 30) / 1000,
"scan-start should be 30 ms"
);
assert_eq!(
TRIM_MIN_GAP_SAMPLES,
(SAMPLE_RATE as usize * 50) / 1000,
"min-gap should be 50 ms"
);
assert_eq!(
TRIM_MAX_DROP_SAMPLES,
(SAMPLE_RATE as usize * 1200) / 1000,
"max-drop should be 1.2 s"
);
}
}

1174
desktop/src-tauri/src/huddle/tts.rs
View File

File diff suppressed because it is too large Load Diff

917
desktop/src-tauri/src/huddle/tts_tests.rs Normal file
View File

@@ -0,0 +1,917 @@
//! Tests for huddle/tts.rs — split into a sibling file to keep tts.rs
//! focused. These exercise the pure helpers and the lifecycle contracts:
//! remote-interrupt frame counting, cancel consumption, the worker idle
//! branch, buffer building, fades, and playback gain.
use super::*;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
// ── Remote interrupt tracker ──────────────────────────────────────────────
//
// Models the per-peer frame counting logic in the recv task of
// relay_api.rs. The contract (must match the production implementation):
//
// - Frame counting is GATED on tts_active — counters only increment
// while TTS is playing.
// - On TTS session start (false→true transition), all counters and
// the window timer are reset. Prevents stale pre-playback speech
// from tripping a cancel.
// - In production, counting happens after successful Opus decode
// (Ok(n) if n > 0). We can't model decode in unit tests, so
// on_frame() represents a successfully-decoded audio frame.
// - Each remote peer has an independent frame counter.
// - Counters use saturating_add (overflow-safe).
// - When a peer's counter crosses REMOTE_SPEECH_THRESHOLD, set
// tts_cancel = true.
// - Counters reset on the 500ms window (Instant-based in production,
// on_tick() in tests — logically equivalent).
// - Uses Acquire for tts_active reads, Release for tts_cancel writes.
//
use crate::huddle::relay_api::REMOTE_SPEECH_THRESHOLD;
/// Test-side model of the per-peer frame counting logic in the recv task.
struct RemoteInterruptTracker {
frame_counts: HashMap<u8, u16>,
tts_active: Arc<AtomicBool>,
tts_cancel: Arc<AtomicBool>,
/// Tracks the previous tts_active state to detect false→true transitions.
/// Mirrors `tts_was_active` in the production recv task.
tts_was_active: bool,
}
impl RemoteInterruptTracker {
fn new(tts_active: Arc<AtomicBool>, tts_cancel: Arc<AtomicBool>) -> Self {
Self {
frame_counts: HashMap::new(),
tts_active,
tts_cancel,
tts_was_active: false,
}
}
/// Called when a successfully-decoded audio frame arrives from a
/// remote peer. Mirrors the production logic in relay_api.rs:
/// 1. Check tts_active — skip if inactive
/// 2. On false→true transition, clear all counters (new TTS session)
/// 3. Increment peer counter (saturating)
/// 4. Fire cancel if threshold crossed
fn on_frame(&mut self, peer_idx: u8) {
let tts_now = self.tts_active.load(Ordering::Acquire);
// Detect TTS session start — clear stale counters.
if tts_now && !self.tts_was_active {
self.frame_counts.clear();
}
self.tts_was_active = tts_now;
if !tts_now {
return; // Not counting while TTS is inactive.
}
let count = self.frame_counts.entry(peer_idx).or_insert(0);
*count = count.saturating_add(1);
if *count >= REMOTE_SPEECH_THRESHOLD {
self.tts_cancel.store(true, Ordering::Release);
}
}
/// Called on the 500ms window boundary — resets all frame counters.
/// In production this is Instant-based (starvation-proof); in tests
/// we call it explicitly since there's no async event loop.
fn on_tick(&mut self) {
self.frame_counts.clear();
}
fn count_for(&self, peer_idx: u8) -> u16 {
*self.frame_counts.get(&peer_idx).unwrap_or(&0)
}
}
/// Simulate the TTS worker's cancel-handling logic (from handle_cancel_or_shutdown).
/// Returns true if cancel was processed (mirrors the real function's return value).
fn simulate_cancel_consumption(
cancel: &AtomicBool,
shutdown: &AtomicBool,
tts_active: &AtomicBool,
text_rx: &mpsc::Receiver<String>,
) -> bool {
if shutdown.load(Ordering::Acquire) {
tts_active.store(false, Ordering::Release);
return true;
}
if cancel.load(Ordering::Acquire) {
while text_rx.try_recv().is_ok() {}
cancel.store(false, Ordering::Release);
tts_active.store(false, Ordering::Release);
return true;
}
false
}
/// Simulate the drained-player check in the TTS worker's main loop. It
/// runs in TWO places — the recv-timeout (idle) arm, and on item receipt
/// before synthesis begins: `tts_active` is released and the lead-in
/// re-armed only when the player has fully drained AND audio was actually
/// queued since the last idle period. Must match the production logic in
/// `tts_worker`.
fn simulate_idle_check(player_empty: bool, first_append: &mut bool, tts_active: &AtomicBool) {
if player_empty && !*first_append {
tts_active.store(false, Ordering::Release);
*first_append = true;
}
}
// ── Threshold tests ───────────────────────────────────────────────────────
/// Real speech above threshold during TTS → cancel fires.
#[test]
fn speech_above_threshold_sets_cancel() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Send exactly REMOTE_SPEECH_THRESHOLD frames from peer 1.
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
}
assert!(
tts_cancel.load(Ordering::Acquire),
"tts_cancel should be true after {} frames from a peer during active TTS",
REMOTE_SPEECH_THRESHOLD,
);
}
/// DTX comfort noise below threshold during TTS → cancel does NOT fire.
#[test]
fn comfort_noise_below_threshold_no_cancel() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// DTX comfort noise: ~2-3 frames per 500ms tick. Send fewer than threshold.
for _ in 0..(REMOTE_SPEECH_THRESHOLD - 1) {
tracker.on_frame(1);
}
assert!(
!tts_cancel.load(Ordering::Acquire),
"tts_cancel should remain false with only {} frames (below threshold of {})",
REMOTE_SPEECH_THRESHOLD - 1,
REMOTE_SPEECH_THRESHOLD,
);
}
/// Frames arrive while tts_active=false → no cancel regardless of count.
#[test]
fn frames_without_tts_active_no_cancel() {
let tts_active = Arc::new(AtomicBool::new(false));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Send many frames — well above threshold.
for _ in 0..50 {
tracker.on_frame(1);
}
assert!(
!tts_cancel.load(Ordering::Acquire),
"tts_cancel should remain false when TTS is not active",
);
}
/// Counter reset on tick → peer must re-accumulate frames to trigger cancel.
#[test]
fn tick_resets_counters() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Send frames just below threshold.
for _ in 0..(REMOTE_SPEECH_THRESHOLD - 1) {
tracker.on_frame(1);
}
assert!(!tts_cancel.load(Ordering::Acquire));
// Tick resets counters.
tracker.on_tick();
assert_eq!(tracker.count_for(1), 0, "counter should be zero after tick");
// Send same number again — still below threshold because counter was reset.
for _ in 0..(REMOTE_SPEECH_THRESHOLD - 1) {
tracker.on_frame(1);
}
assert!(
!tts_cancel.load(Ordering::Acquire),
"cancel should not fire: counter was reset by tick, frames still below threshold",
);
}
/// Per-peer isolation: peer A's silence does not reset peer B's accumulation.
#[test]
fn per_peer_counters_are_independent() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Peer 1 sends frames below threshold (DTX comfort noise).
for _ in 0..2 {
tracker.on_frame(1);
}
// Peer 2 sends frames above threshold (real speech).
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(2);
}
assert!(
tts_cancel.load(Ordering::Acquire),
"peer 2 should trigger cancel independently of peer 1's low count",
);
assert_eq!(
tracker.count_for(1),
2,
"peer 1 counter should be untouched"
);
assert_eq!(
tracker.count_for(2),
REMOTE_SPEECH_THRESHOLD,
"peer 2 counter should be at threshold",
);
}
/// Multiple peers both above threshold → cancel fires (idempotent).
#[test]
fn multiple_peers_above_threshold_idempotent() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Both peers send above threshold.
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
tracker.on_frame(2);
}
assert!(
tts_cancel.load(Ordering::Acquire),
"cancel should be set when multiple peers exceed threshold",
);
// tts_active should still be true — only the TTS worker resets it.
assert!(
tts_active.load(Ordering::Acquire),
"tts_active should remain true (only TTS worker clears it)",
);
}
/// tts_cancel already true → setting again is harmless (AtomicBool store).
#[test]
fn cancel_already_true_is_harmless() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(true)); // Already cancelled.
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
}
assert!(
tts_cancel.load(Ordering::Acquire),
"tts_cancel should still be true (idempotent store)",
);
}
// ── Regression: local-only interrupt still works ──────────────────────────
/// The existing local barge-in path (STT detects speech → sets tts_cancel)
/// must continue to work independently of remote frame counting.
#[test]
fn local_barge_in_still_works_without_remote_frames() {
let _tts_active = AtomicBool::new(true);
let tts_cancel = AtomicBool::new(false);
// Simulate local STT barge-in (stt.rs after BARGE_IN_DEBOUNCE_FRAMES).
tts_cancel.store(true, Ordering::Release);
assert!(
tts_cancel.load(Ordering::Acquire),
"local barge-in should set tts_cancel",
);
}
// ── Cancel consumption tests (TTS worker side) ────────────────────────────
/// TTS worker correctly resets both tts_cancel and tts_active after cancel.
#[test]
fn cancel_consumption_resets_flags() {
let cancel = AtomicBool::new(true);
let shutdown = AtomicBool::new(false);
let tts_active = AtomicBool::new(true);
let (_tx, rx) = mpsc::channel::<String>();
let handled = simulate_cancel_consumption(&cancel, &shutdown, &tts_active, &rx);
assert!(handled, "should return true when cancel is set");
assert!(
!cancel.load(Ordering::Acquire),
"cancel should be reset to false after consumption",
);
assert!(
!tts_active.load(Ordering::Acquire),
"tts_active should be reset to false after cancel",
);
}
/// TTS worker drains the text queue on cancel.
#[test]
fn cancel_consumption_drains_queue() {
let cancel = AtomicBool::new(true);
let shutdown = AtomicBool::new(false);
let tts_active = AtomicBool::new(true);
let (tx, rx) = mpsc::channel::<String>();
tx.send("sentence one".to_string()).unwrap();
tx.send("sentence two".to_string()).unwrap();
tx.send("sentence three".to_string()).unwrap();
let handled = simulate_cancel_consumption(&cancel, &shutdown, &tts_active, &rx);
assert!(handled);
assert!(
rx.try_recv().is_err(),
"text queue should be drained after cancel",
);
}
/// No cancel, no shutdown → TTS worker continues (returns false).
#[test]
fn no_cancel_no_shutdown_returns_false() {
let cancel = AtomicBool::new(false);
let shutdown = AtomicBool::new(false);
let tts_active = AtomicBool::new(true);
let (_tx, rx) = mpsc::channel::<String>();
let handled = simulate_cancel_consumption(&cancel, &shutdown, &tts_active, &rx);
assert!(
!handled,
"should return false when neither cancel nor shutdown is set",
);
assert!(
tts_active.load(Ordering::Acquire),
"tts_active should remain true",
);
}
/// Shutdown takes priority over cancel — cancel flag is NOT reset.
#[test]
fn shutdown_takes_priority_over_cancel() {
let cancel = AtomicBool::new(true);
let shutdown = AtomicBool::new(true);
let tts_active = AtomicBool::new(true);
let (_tx, rx) = mpsc::channel::<String>();
let handled = simulate_cancel_consumption(&cancel, &shutdown, &tts_active, &rx);
assert!(handled, "should return true on shutdown");
assert!(
!tts_active.load(Ordering::Acquire),
"tts_active should be false after shutdown",
);
assert!(
cancel.load(Ordering::Acquire),
"cancel should remain true (shutdown path doesn't reset it)",
);
}
// ── Full lifecycle tests ──────────────────────────────────────────────────
/// Idle branch: nothing queued + player drained + audio was played →
/// release `tts_active` and re-arm the lead-in for the next utterance.
#[test]
fn idle_check_releases_mic_gate_after_drain() {
let tts_active = AtomicBool::new(true);
let mut first_append = false; // audio was appended this utterance
simulate_idle_check(true, &mut first_append, &tts_active);
assert!(
!tts_active.load(Ordering::Acquire),
"tts_active should be released once playback drains",
);
assert!(
first_append,
"lead-in should be re-armed for next utterance"
);
}
/// Idle branch: nothing queued but audio is STILL DRAINING → `tts_active`
/// stays set. This is the cross-item pipelining contract: the per-item
/// drain barrier is gone, so the mic gate must hold across the gap
/// between item N finishing synthesis and its audio finishing playback.
#[test]
fn idle_check_holds_mic_gate_while_audio_drains() {
let tts_active = AtomicBool::new(true);
let mut first_append = false;
simulate_idle_check(false, &mut first_append, &tts_active);
assert!(
tts_active.load(Ordering::Acquire),
"tts_active must stay set while queued audio is still playing",
);
assert!(!first_append, "lead-in must not re-arm mid-playback");
}
/// Idle branch: player empty but nothing was ever appended (e.g. worker
/// just started, or every item preprocessed to empty) → no spurious
/// store, lead-in stays armed.
#[test]
fn idle_check_noop_when_nothing_was_queued() {
let tts_active = AtomicBool::new(false);
let mut first_append = true;
simulate_idle_check(true, &mut first_append, &tts_active);
assert!(!tts_active.load(Ordering::Acquire));
assert!(first_append, "lead-in should remain armed");
}
/// On-receipt check (PR #997 review blocker): item N's audio drains, then
/// item N+1 arrives BEFORE the recv timeout fires. The same drained-player
/// check must run on receipt so `tts_active` is released before the
/// synthesis pass — otherwise STT discards human speech as "echo" for the
/// whole synthesis window while the agent is actually silent.
#[test]
fn on_receipt_check_releases_mic_gate_before_synthesis() {
let tts_active = AtomicBool::new(true);
let mut first_append = false; // item N appended audio, now fully drained
// Item N+1 received with the player already empty: release + re-arm.
simulate_idle_check(true, &mut first_append, &tts_active);
assert!(
!tts_active.load(Ordering::Acquire),
"tts_active must be released before synthesizing the next item \
when playback has already drained",
);
assert!(
first_append,
"lead-in must re-arm so the next utterance gets a fresh cushion"
);
}
// ── Barge-in monitor ──────────────────────────────────────────────────────
//
// Models one tick of the tts-barge-in-monitor thread in `tts_worker`. The
// contract (must match production):
// - cancel set → take `player_ops`, RE-CHECK cancel under the lock; if
// still set: silence player + release tts_active, flag NOT consumed
// (the worker owns consumption: queue drain + lead-in reset)
// - cancel cleared by the time the lock is held → no-op (stale branch)
// - cancel clear → no-op
// Not consuming the flag is what makes the monitor idempotent across ticks
// and closes the race where the worker appends a sentence after the
// monitor's clear but before consuming the flag. The under-lock re-check
// closes the converse race (PR #997 review blocker): worker consumes the
// cancel and appends a fresh post-cancel utterance between the monitor's
// initial load and its clear — the stale branch must not delete that audio.
/// Test-side model of one monitor tick. `player_cleared` stands in for the
/// `clear()+play()` pair on the real Player; `player_ops` is the mutex
/// serializing monitor clears with worker player mutations.
fn simulate_monitor_tick(
cancel: &AtomicBool,
tts_active: &AtomicBool,
player_ops: &Mutex<()>,
player_cleared: &mut bool,
) {
if cancel.load(Ordering::Acquire) {
let _ops = player_ops.lock().unwrap();
if cancel.load(Ordering::Acquire) {
*player_cleared = true;
tts_active.store(false, Ordering::Release);
}
}
}
/// Cancel set → monitor silences playback and releases the mic gate, but
/// leaves the flag for the worker to consume.
#[test]
fn monitor_tick_silences_and_releases_without_consuming_cancel() {
let cancel = AtomicBool::new(true);
let tts_active = AtomicBool::new(true);
let player_ops = Mutex::new(());
let mut player_cleared = false;
simulate_monitor_tick(&cancel, &tts_active, &player_ops, &mut player_cleared);
assert!(player_cleared, "monitor must silence in-flight audio");
assert!(
!tts_active.load(Ordering::Acquire),
"monitor must release the mic gate immediately on barge-in",
);
assert!(
cancel.load(Ordering::Acquire),
"monitor must NOT consume the cancel flag — the worker owns \
queue drain and lead-in reset",
);
}
/// No cancel → monitor is a pure no-op.
#[test]
fn monitor_tick_noop_without_cancel() {
let cancel = AtomicBool::new(false);
let tts_active = AtomicBool::new(true);
let player_ops = Mutex::new(());
let mut player_cleared = false;
simulate_monitor_tick(&cancel, &tts_active, &player_ops, &mut player_cleared);
assert!(!player_cleared);
assert!(
tts_active.load(Ordering::Acquire),
"monitor must not touch the mic gate while no barge-in is pending",
);
}
/// Stale-branch race (PR #997 review blocker): monitor observes
/// `cancel == true`, then the worker — under `player_ops` — consumes the
/// cancel and appends a fresh post-cancel utterance before the monitor
/// reaches its clear. The monitor's under-lock re-check must see
/// `cancel == false` and no-op, leaving the fresh audio and its mic gate
/// intact.
#[test]
fn monitor_stale_cancel_branch_must_not_clear_fresh_audio() {
let cancel = AtomicBool::new(true);
let tts_active = AtomicBool::new(true);
let player_ops = Mutex::new(());
let mut player_cleared = false;
// Monitor's initial (pre-lock) load observes the cancel…
let stale_observation = cancel.load(Ordering::Acquire);
assert!(stale_observation);
// …then the worker wins the lock: consumes the cancel, appends a fresh
// utterance, sets tts_active (mirrors handle_cancel_or_shutdown +
// the locked append in the sentence loop).
{
let _ops = player_ops.lock().unwrap();
cancel.store(false, Ordering::Release);
tts_active.store(true, Ordering::Release);
}
// Monitor resumes from its stale branch — the re-check under the lock
// must turn it into a no-op.
if stale_observation {
let _ops = player_ops.lock().unwrap();
if cancel.load(Ordering::Acquire) {
player_cleared = true;
tts_active.store(false, Ordering::Release);
}
}
assert!(
!player_cleared,
"stale monitor branch must not clear a fresh post-cancel utterance",
);
assert!(
tts_active.load(Ordering::Acquire),
"stale monitor branch must not release the mic gate while fresh \
audio is playing",
);
}
/// Full cycle: remote speech → cancel → TTS consumption → new TTS → cancel again.
/// Validates the cancel mechanism is reusable across TTS sessions.
/// The false→true transition on tts_active auto-clears counters — no
/// explicit on_tick() needed between sessions.
#[test]
fn full_cancel_cycle_is_reusable() {
let tts_active = Arc::new(AtomicBool::new(false));
let tts_cancel = Arc::new(AtomicBool::new(false));
let shutdown = AtomicBool::new(false);
let (_tx, rx) = mpsc::channel::<String>();
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Cycle 1: TTS starts, remote speech triggers cancel, TTS consumes.
tts_active.store(true, Ordering::Release);
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
}
assert!(tts_cancel.load(Ordering::Acquire));
let handled = simulate_cancel_consumption(&tts_cancel, &shutdown, &tts_active, &rx);
assert!(handled);
assert!(!tts_cancel.load(Ordering::Acquire));
assert!(!tts_active.load(Ordering::Acquire));
// No on_tick() needed — the false→true transition auto-clears counters.
// Cycle 2: New TTS starts, another remote speech triggers cancel.
tts_active.store(true, Ordering::Release);
for _ in 0..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
}
assert!(tts_cancel.load(Ordering::Acquire));
let handled = simulate_cancel_consumption(&tts_cancel, &shutdown, &tts_active, &rx);
assert!(handled);
assert!(!tts_cancel.load(Ordering::Acquire));
assert!(!tts_active.load(Ordering::Acquire));
}
/// TTS session transition (false→true) clears stale counters.
/// Prevents pre-existing speech from tripping a cancel on TTS restart.
#[test]
fn tts_session_transition_clears_counters() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Accumulate frames just below threshold during first TTS session.
for _ in 0..(REMOTE_SPEECH_THRESHOLD - 1) {
tracker.on_frame(1);
}
assert_eq!(tracker.count_for(1), REMOTE_SPEECH_THRESHOLD - 1);
assert!(!tts_cancel.load(Ordering::Acquire));
// TTS stops (cancel consumed).
tts_active.store(false, Ordering::Release);
// Send a frame while inactive — triggers tts_was_active transition tracking.
tracker.on_frame(1);
// TTS restarts — false→true transition should clear counters.
tts_active.store(true, Ordering::Release);
tracker.on_frame(1); // First frame of new session triggers clear + count.
assert_eq!(
tracker.count_for(1),
1,
"counter should be 1 (cleared on session transition, then incremented)",
);
// Need full threshold again from scratch.
for _ in 1..REMOTE_SPEECH_THRESHOLD {
tracker.on_frame(1);
}
assert!(
tts_cancel.load(Ordering::Acquire),
"cancel should fire after full threshold in new session",
);
}
/// Concurrent remote cancel + local barge-in → both safe, one cancel processed.
#[test]
fn concurrent_remote_and_local_cancel() {
let tts_active = Arc::new(AtomicBool::new(true));
let tts_cancel = Arc::new(AtomicBool::new(false));
// Remote path: frame counting above threshold.
let cancel_remote = Arc::clone(&tts_cancel);
let active_remote = Arc::clone(&tts_active);
let remote = std::thread::spawn(move || {
// Simulate threshold crossing — directly set cancel (the tracker
// would do this after REMOTE_SPEECH_THRESHOLD frames).
if active_remote.load(Ordering::Acquire) {
cancel_remote.store(true, Ordering::Release);
}
});
// Local path: STT barge-in.
let cancel_local = Arc::clone(&tts_cancel);
let local = std::thread::spawn(move || {
cancel_local.store(true, Ordering::Release);
});
remote.join().unwrap();
local.join().unwrap();
assert!(
tts_cancel.load(Ordering::Acquire),
"tts_cancel should be true regardless of which path set it",
);
}
/// Frames while TTS is inactive are NOT counted. The peer must accumulate
/// the full threshold AFTER tts_active becomes true.
#[test]
fn frames_while_tts_inactive_are_not_counted() {
let tts_active = Arc::new(AtomicBool::new(false));
let tts_cancel = Arc::new(AtomicBool::new(false));
let mut tracker = RemoteInterruptTracker::new(Arc::clone(&tts_active), Arc::clone(&tts_cancel));
// Send frames while TTS is inactive — should be ignored entirely.
for _ in 0..50 {
tracker.on_frame(1);
}
assert_eq!(
tracker.count_for(1),
0,
"no frames should accumulate while TTS inactive"
);
assert!(!tts_cancel.load(Ordering::Acquire));
// TTS starts. Peer must now accumulate from zero.
tts_active.store(true, Ordering::Release);
// Send frames below threshold — not enough yet.
for _ in 0..(REMOTE_SPEECH_THRESHOLD - 1) {
tracker.on_frame(1);
}
assert!(
!tts_cancel.load(Ordering::Acquire),
"cancel should not fire: only {} frames since TTS became active",
REMOTE_SPEECH_THRESHOLD - 1,
);
// One more frame crosses the threshold.
tracker.on_frame(1);
assert!(
tts_cancel.load(Ordering::Acquire),
"cancel should fire after full threshold accumulated post-activation",
);
}
// ── apply_fade_out tests ──────────────────────────────────────────────────
/// The fade-out half of the old `apply_fades`: last sample is silenced
/// and the ramp is monotonic. Mid-buffer must be untouched.
#[test]
fn apply_fade_out_short_buffer() {
let mut samples = vec![1.0f32; 10];
apply_fade_out(&mut samples);
assert_eq!(samples[9], 0.0, "last sample should be silenced");
assert!(samples[5] > 0.5, "mid-buffer should be near-untouched");
}
/// REGRESSION (2026-05-18): the *first* samples must NOT be attenuated.
/// An earlier `apply_fades` symmetrically faded in over 8 ms which
/// swallowed the consonant onset of every sentence.
/// Lock in: samples[0..FADE_OUT_SAMPLES] are byte-equal to input.
#[test]
fn apply_fade_out_does_not_touch_leading_samples() {
// Input long enough that fade window doesn't overlap (≫ 2× fade).
let n = FADE_OUT_SAMPLES * 4;
let input: Vec<f32> = (0..n).map(|i| 0.5 + (i as f32) * 1e-4).collect();
let mut samples = input.clone();
apply_fade_out(&mut samples);
for i in 0..FADE_OUT_SAMPLES {
assert_eq!(
samples[i], input[i],
"leading sample {i} must not be attenuated (was {} → {})",
input[i], samples[i]
);
}
// And the trailing fade still works.
assert_eq!(samples[n - 1], 0.0);
}
#[test]
fn apply_fade_out_empty_buffer() {
let mut samples: Vec<f32> = vec![];
apply_fade_out(&mut samples);
assert!(samples.is_empty());
}
#[test]
fn apply_fade_out_single_sample() {
// fade = min(FADE_OUT_SAMPLES, len/2) = 0, so nothing changes.
let mut samples = vec![1.0f32];
apply_fade_out(&mut samples);
assert_eq!(samples[0], 1.0);
}
/// Sanity-check the per-sentence cushion length: 20 ms at 24 kHz must
/// land at exactly 480 samples. This is a const computation, so the
/// real value of this test is documenting *why* 20 ms was chosen — it
/// covers a typical CoreAudio buffer turnover (256–1024 samples)
/// without being audible as user-facing latency.
#[test]
fn sentence_lead_in_is_sane() {
assert_eq!(SENTENCE_LEAD_IN_SAMPLES, 480, "20 ms × 24 kHz");
}
// ── build_sentence_append_buffer tests ───────────────────────────────────
/// REGRESSION: every chunk needs an onset cushion; synthesized chunks
/// can start with speech energy within the first millisecond.
#[test]
fn lead_in_pad_is_present_for_every_sentence_chunk() {
const SENTENCE_AUDIO_LEN: usize = 1000;
const SILENCE_BUF_LEN: usize = 2400; // 100 ms at 24 kHz, like production
const N_SENTENCES: usize = 5;
let mut first = true;
for _ in 0..N_SENTENCES {
let buf = build_sentence_append_buffer(
&mut first,
vec![0.5_f32; SENTENCE_AUDIO_LEN],
SILENCE_BUF_LEN,
);
assert_eq!(buf.len(), SENTENCE_AUDIO_LEN + SILENCE_BUF_LEN);
assert!(
buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0),
"lead-in pad must be pure silence"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN]
.iter()
.all(|&s| s == 0.5),
"sentence audio must immediately follow the lead-in"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN..]
.iter()
.all(|&s| s == 0.0),
"trailing gap must be pure silence"
);
}
assert!(!first, "first_append flag must be cleared after first call");
}
/// `first_append` still flips on the first call for `tts_active` gating.
#[test]
fn build_sentence_append_buffer_flips_first_append() {
let mut first = true;
let _ = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(!first, "first call must flip the flag");
// Subsequent call: still has a per-sentence lead-in, flag stays false.
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(!first);
}
/// Leading silence is exactly the lead-in; no pre-audio gap is double-counted.
#[test]
fn first_sentence_leading_silence_is_exactly_lead_in() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert_eq!(buf[SENTENCE_LEAD_IN_SAMPLES], 0.5);
}
/// Tail silence plus the next lead-in preserves the 100 ms sentence gap.
#[test]
fn sentence_gap_budget_is_preserved() {
let mut first = true;
let silence_buf_len = 2400;
let first_buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len);
let second_buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len);
let first_tail = &first_buf[SENTENCE_LEAD_IN_SAMPLES + 100..];
let second_lead = &second_buf[..SENTENCE_LEAD_IN_SAMPLES];
assert_eq!(first_tail.len(), silence_buf_len - SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(second_lead.len(), SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(first_tail.len() + second_lead.len(), silence_buf_len);
}
/// Regression guard: one contiguous rodio source per synthesized sentence.
#[test]
fn sentence_append_buffer_is_one_contiguous_source() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400);
assert_eq!(buf.len(), 2400 + 100);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + 100]
.iter()
.all(|&s| s == 0.5)
);
}
// ── clamp_to_full_scale tests ─────────────────────────────────────────────
/// In-range speech audio passes through bit-exact — no gain is applied.
/// (Pocket output is already at speech level; see the fn doc-comment for
/// the history of the two gain-stage regressions this replaced.)
#[test]
fn clamp_to_full_scale_passes_speech_audio_unchanged() {
let input = vec![0.42_f32, -0.97, 0.076, 0.0];
let out = clamp_to_full_scale(input.clone());
assert_eq!(out, input);
}
/// Outlier transients beyond full scale are hard-clamped to ±1.0 rather
/// than wrapping.
#[test]
fn clamp_to_full_scale_clamps_outliers() {
let input = vec![1.5_f32, -2.0, 0.5];
let out = clamp_to_full_scale(input);
assert_eq!(out, vec![1.0, -1.0, 0.5]);
}
/// Empty input round-trips to empty output.
#[test]
fn clamp_to_full_scale_empty_buffer() {
let out = clamp_to_full_scale(Vec::new());
assert!(out.is_empty());
}