pipecrab_vad/stage.rs
1//! [`VadStage`]: the [`VoiceActivityDetector`] gate. Audio in; speech-only audio
2//! out, bracketed by well-formed speech edges.
3
4use std::collections::VecDeque;
5use std::sync::Mutex;
6use std::time::Duration;
7
8use async_trait::async_trait;
9use pipecrab_core::{
10 AudioChunk, AudioFormat, DataFrame, Decision, Direction, Processor, SystemFrame,
11};
12use pipecrab_runtime::{Outbound, Stage, StageError};
13
14use crate::{VadError, VadEvent, VoiceActivityDetector};
15
16/// Adapts any [`VoiceActivityDetector`] into a pipeline [`Stage`] — as a
17/// **gate**, not a tap.
18///
19/// The stage runs the detector over every conforming audio chunk and emits, on
20/// the data lane, an utterance's audio **bracketed by its edges**: a
21/// [`SpeechStarted`](DataFrame::SpeechStarted), then the utterance's chunks
22/// (pre-roll included), then a [`SpeechStopped`](DataFrame::SpeechStopped).
23/// While idle it emits *nothing* — silence dies here.
24///
25/// # ⚠️ Contract inversion
26///
27/// The older lane-discipline design made VAD a tap: the audio flowed through
28/// untouched and the edge was emitted **after** the chunk that triggered it.
29/// The gate **inverts** that. Downstream of `VadStage`, **edges bracket the
30/// utterance audio**: `SpeechStarted` *precedes* every utterance chunk (pre-roll
31/// included) and `SpeechStopped` *follows* the last one. A stateless downstream
32/// stage can therefore open its utterance on the edge alone — the onset audio is
33/// already gated in behind it.
34///
35/// # The pre-roll ring
36///
37/// A detector only declares speech *started* after enough speech has accrued
38/// ([`DebounceConfig::start_windows`](crate::DebounceConfig)), so real onset
39/// audio has already passed by the time the edge fires. The gate owns a
40/// duration-bounded pre-roll ring: while idle it stashes each incoming chunk,
41/// and on onset it flushes the ring — in arrival order — ahead of the triggering
42/// chunk, so the utterance's first syllables survive. The budget is
43/// [`GateConfig::preroll`] (default 300 ms).
44///
45/// # Topology commitment
46///
47/// Because the gate drops silence, any future consumer of *continuous* raw audio
48/// (recording, a level meter, an AEC reference) must sit **upstream** of
49/// `VadStage`. Nothing downstream consumes silence today, so this constrains
50/// nothing yet, but it is a standing commitment (see `ARCHITECTURE.md`).
51///
52/// # State and cancellation
53///
54/// The detector's own edge state lives in the detector (reset via its control
55/// call). The stage owns only the gate — the current mode and the ring —
56/// behind a [`Mutex`], because [`perform`](Stage::perform) is `&self`. All gate
57/// mutation happens in one synchronous critical section *after* the awaited
58/// [`process`](VoiceActivityDetector::process); the sends happen after the lock
59/// is released. A `perform` dropped mid-send loses unsent onset audio but can
60/// never resurrect a stale chunk into a later utterance.
61pub struct VadStage<V: VoiceActivityDetector> {
62 detector: V,
63 /// The one format the detector accepts, cached from
64 /// [`input_format`](VoiceActivityDetector::input_format) in [`new`](Self::new).
65 expected: AudioFormat,
66 /// Mode + ring. Locked only in post-await critical sections.
67 gate: Mutex<Gate>,
68}
69
70/// Tuning for [`VadStage`]'s pre-roll ring.
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub struct GateConfig {
73 /// How much onset audio to retain in the pre-roll ring. Larger keeps more of
74 /// the utterance's first syllables at the cost of buffering; the default is
75 /// 300 ms.
76 pub preroll: Duration,
77}
78
79impl Default for GateConfig {
80 fn default() -> Self {
81 Self {
82 preroll: Duration::from_millis(300),
83 }
84 }
85}
86
87/// The gate's mutable state: whether we are passing speech through, and the
88/// pre-roll ring that accumulates while idle.
89struct Gate {
90 mode: Mode,
91 ring: PrerollRing,
92}
93
94/// Which side of the gate we are on.
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96enum Mode {
97 /// No open utterance: incoming audio accumulates in the ring, nothing is
98 /// emitted.
99 Idle,
100 /// An utterance is open: incoming audio flows through as [`DataFrame::Audio`].
101 Speech,
102}
103
104/// A duration-bounded FIFO of audio chunks: the pre-roll buffer that captures an
105/// utterance's onset — the audio that arrives *before* the detector's
106/// `SpeechStarted` edge — so the first syllables are not clipped.
107///
108/// While the gate is idle it accumulates chunks, evicting the oldest whole
109/// chunks once the total buffered duration exceeds `budget`. Chunks vary in
110/// size, so eviction works in whole chunks rather than samples. Admission is
111/// format-fatal upstream (the stage rejects a mismatch before it reaches here),
112/// so the ring is uniform-format by construction — it needs no clear-on-change
113/// logic of its own.
114struct PrerollRing {
115 /// The maximum total duration to retain.
116 budget: Duration,
117 /// Buffered chunks, oldest at the front.
118 chunks: VecDeque<AudioChunk>,
119}
120
121impl PrerollRing {
122 fn new(budget: Duration) -> Self {
123 Self {
124 budget,
125 chunks: VecDeque::new(),
126 }
127 }
128
129 /// The total buffered duration. Recomputed from the chunks (the ring is
130 /// small, ~budget/window) so accounting never drifts.
131 fn total(&self) -> Duration {
132 self.chunks.iter().map(chunk_duration).sum()
133 }
134
135 /// Push a chunk, evicting oldest whole chunks to honour the duration budget.
136 fn push(&mut self, chunk: AudioChunk) {
137 self.chunks.push_back(chunk);
138 // Evict oldest whole chunks until we fit the budget, but always keep the
139 // most recent chunk: a lone chunk longer than the whole budget is still
140 // the freshest onset audio, and dropping it would clip the utterance.
141 while self.chunks.len() > 1 && self.total() > self.budget {
142 self.chunks.pop_front();
143 }
144 }
145
146 /// Remove and return every buffered chunk in arrival order.
147 fn take(&mut self) -> Vec<AudioChunk> {
148 self.chunks.drain(..).collect()
149 }
150
151 /// Discard every buffered chunk.
152 fn clear(&mut self) {
153 self.chunks.clear();
154 }
155}
156
157/// The wall-clock duration of one audio chunk: interleaved frames over the
158/// sample rate. A malformed format (zero rate) yields zero rather than dividing
159/// by it.
160fn chunk_duration(chunk: &AudioChunk) -> Duration {
161 let channels = chunk.format.channels.max(1) as u64;
162 let rate = chunk.format.sample_rate as u64;
163 if rate == 0 {
164 return Duration::ZERO;
165 }
166 let frames = chunk.samples.len() as u64 / channels;
167 // Integer nanoseconds keep the budget accounting exact and drift-free.
168 Duration::from_nanos(frames * 1_000_000_000 / rate)
169}
170
171impl<V: VoiceActivityDetector> VadStage<V> {
172 /// Wrap `detector` as a gate with the default [`GateConfig`].
173 pub fn new(detector: V) -> Self {
174 Self::with_config(detector, GateConfig::default())
175 }
176
177 /// Wrap `detector` as a gate with an explicit [`GateConfig`].
178 pub fn with_config(detector: V, config: GateConfig) -> Self {
179 let expected = detector.input_format();
180 Self {
181 detector,
182 expected,
183 gate: Mutex::new(Gate {
184 mode: Mode::Idle,
185 ring: PrerollRing::new(config.preroll),
186 }),
187 }
188 }
189}
190
191/// One step the gate asks [`perform`](Stage::perform) to carry out:
192/// [`VadStage`]'s [`Processor::Effect`]. Its contents are private — only the
193/// stage constructs one — so the effect vocabulary stays opaque to callers.
194pub struct VadEffect(Effect);
195
196enum Effect {
197 /// Run detection over this conforming chunk and drive the gate.
198 Detect(AudioChunk),
199 /// The chunk's format did not match the detector's; fail fatally.
200 RejectFormat { got: AudioFormat },
201}
202
203impl<V: VoiceActivityDetector> Processor for VadStage<V> {
204 type Effect = VadEffect;
205
206 fn decide_data(&mut self, frame: &DataFrame) -> Decision<VadEffect> {
207 match frame {
208 // Format-fatal admission: a mismatch is rejected before any audio
209 // reaches the engine (the engine cannot detect rate from samples).
210 DataFrame::Audio(chunk) if chunk.format != self.expected => {
211 Decision::drop().emit(VadEffect(Effect::RejectFormat { got: chunk.format }))
212 }
213 // Conforming audio: drop it — the gate owns forwarding now — and let
214 // `perform` decide whether it is gated through or stashed in the ring.
215 // The chunk is Arc-backed, so this clone is a refcount bump.
216 DataFrame::Audio(chunk) => {
217 Decision::drop().emit(VadEffect(Effect::Detect(chunk.clone())))
218 }
219 // Everything else is not ours to inspect.
220 _ => Decision::forward(),
221 }
222 }
223
224 fn decide_system(&mut self, _dir: Direction, frame: &SystemFrame) -> Decision<VadEffect> {
225 match frame {
226 SystemFrame::Interrupt => {
227 // An Interrupt reaching VadStage is a head-injected session
228 // abandon — turn-manager interrupts originate downstream and never
229 // travel up here — so gate/downstream protocol coherence across it
230 // is explicitly not guaranteed. Both sides re-sync to idle: the
231 // gate resets, and SttStage cancels on the same Interrupt.
232 {
233 let mut gate = self.gate.lock().expect("VAD gate mutex poisoned");
234 gate.mode = Mode::Idle;
235 gate.ring.clear();
236 }
237 // Control call: return the detector to its idle, no-speech state.
238 self.detector.reset();
239 Decision::forward()
240 }
241 // Start, Stop, Error, and any future frames: pass through untouched.
242 _ => Decision::forward(),
243 }
244 }
245}
246
247#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
248#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
249impl<V: VoiceActivityDetector> Stage for VadStage<V> {
250 async fn perform(&self, effect: VadEffect, out: &Outbound) -> Result<(), StageError> {
251 let chunk = match effect.0 {
252 Effect::RejectFormat { got } => {
253 return Err(StageError::fatal(format!(
254 "VadStage requires {} Hz/{} ch (declared by the engine); \
255 got {} Hz/{} ch — insert a resample stage upstream or \
256 reconfigure the source",
257 self.expected.sample_rate,
258 self.expected.channels,
259 got.sample_rate,
260 got.channels,
261 )));
262 }
263 Effect::Detect(chunk) => chunk,
264 };
265
266 let events = self.detector.process(chunk.samples.clone()).await?;
267
268 // ONE critical section: take the ring, flip the mode, and build the send
269 // plan atomically. No await inside; the sends happen after the unlock.
270 let plan: Vec<DataFrame> = {
271 let mut gate = self.gate.lock().expect("VAD gate mutex poisoned");
272 let mut plan = Vec::new();
273 let mut sent_chunk = false;
274 for event in &events {
275 match event {
276 VadEvent::SpeechStarted => {
277 // Alternation invariant: a SpeechStarted while already in
278 // speech is a misbehaving detector. Surface it loudly in
279 // debug; in release, re-opening from Idle mode below still
280 // produces coherent, bracketed output.
281 debug_assert!(
282 gate.mode == Mode::Idle,
283 "VAD alternation violated: SpeechStarted while already in speech",
284 );
285 plan.push(DataFrame::SpeechStarted);
286 // The whole onset window survives: the ring's chunks in
287 // arrival order, then the triggering chunk.
288 for pre in gate.ring.take() {
289 plan.push(DataFrame::Audio(pre));
290 }
291 plan.push(DataFrame::Audio(chunk.clone()));
292 sent_chunk = true;
293 gate.mode = Mode::Speech;
294 }
295 VadEvent::SpeechStopped => {
296 debug_assert!(
297 gate.mode == Mode::Speech,
298 "VAD alternation violated: SpeechStopped while not in speech",
299 );
300 // The tail chunk closes the utterance, then the edge —
301 // unless a Started in this same batch already sent it.
302 if !sent_chunk {
303 plan.push(DataFrame::Audio(chunk.clone()));
304 sent_chunk = true;
305 }
306 plan.push(DataFrame::SpeechStopped);
307 gate.mode = Mode::Idle;
308 }
309 }
310 }
311 if events.is_empty() {
312 match gate.mode {
313 // Live speech: the chunk flows straight through.
314 Mode::Speech => plan.push(DataFrame::Audio(chunk.clone())),
315 // Idle silence: accumulate it in the ring, emit nothing.
316 Mode::Idle => gate.ring.push(chunk.clone()),
317 }
318 }
319 plan
320 };
321
322 for frame in plan {
323 // Ignore the send error: it only fires once the sink has gone away
324 // during shutdown, matching the runtime's own forward path.
325 let _ = out.send_data(frame).await;
326 }
327 Ok(())
328 }
329}
330
331impl From<VadError> for StageError {
332 fn from(e: VadError) -> Self {
333 // A failed detection is recoverable: skip this chunk and keep the
334 // pipeline alive. Only the format path (RejectFormat) is fatal.
335 StageError::new(e.to_string())
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342 use std::sync::Arc;
343
344 fn chunk(sample_rate: u32, channels: u16, samples: usize) -> AudioChunk {
345 AudioChunk::new(
346 Arc::from(vec![0.0f32; samples]),
347 AudioFormat::new(sample_rate, channels),
348 )
349 }
350
351 #[test]
352 fn chunk_duration_is_frames_over_sample_rate() {
353 // 16 000 mono samples at 16 kHz is exactly one second.
354 assert_eq!(
355 chunk_duration(&chunk(16_000, 1, 16_000)),
356 Duration::from_secs(1)
357 );
358 // 1 kHz mono makes one sample == one millisecond.
359 assert_eq!(
360 chunk_duration(&chunk(1_000, 1, 250)),
361 Duration::from_millis(250)
362 );
363 }
364
365 #[test]
366 fn chunk_duration_counts_interleaved_frames_not_samples() {
367 // Stereo: 480 interleaved samples is 240 frames, so 240/48k = 5 ms — half
368 // what a naive samples/rate would give.
369 assert_eq!(
370 chunk_duration(&chunk(48_000, 2, 480)),
371 Duration::from_millis(5)
372 );
373 }
374
375 #[test]
376 fn chunk_duration_of_empty_or_degenerate_is_zero() {
377 assert_eq!(chunk_duration(&chunk(16_000, 1, 0)), Duration::ZERO);
378 // A zero sample rate can't yield a duration; guard rather than divide by it.
379 assert_eq!(chunk_duration(&chunk(0, 1, 100)), Duration::ZERO);
380 }
381
382 #[test]
383 fn preroll_evicts_by_duration_keeping_the_most_recent() {
384 // 1000 Hz mono makes 1 sample == 1 ms, so a 100 ms budget holds ~100 samples.
385 let mut ring = PrerollRing::new(Duration::from_millis(100));
386 // +20 -> [20] (20 ms)
387 // +50 -> [20,50] (70 ms)
388 // +40 (110) -> evict 20 -> [50,40] (90 ms)
389 // +30 (120) -> evict 50 -> [40,30] (70 ms)
390 for n in [20usize, 50, 40, 30] {
391 ring.push(chunk(1000, 1, n));
392 }
393 let survivors: Vec<usize> = ring.take().iter().map(|c| c.samples.len()).collect();
394 assert_eq!(
395 survivors,
396 vec![40, 30],
397 "only the last two chunks survive, in arrival order"
398 );
399 }
400
401 #[test]
402 fn preroll_keeps_a_lone_oversized_chunk() {
403 // A single chunk longer than the whole budget is still the freshest onset
404 // audio: keep it rather than clip the utterance.
405 let mut ring = PrerollRing::new(Duration::from_millis(10));
406 ring.push(chunk(1000, 1, 500)); // 500 ms >> 10 ms budget
407 let survivors: Vec<usize> = ring.take().iter().map(|c| c.samples.len()).collect();
408 assert_eq!(
409 survivors,
410 vec![500],
411 "the most-recent chunk is never evicted"
412 );
413 }
414}