use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use crate::error::Result;
use crate::processor::frame::{Frame, FrameKind};
use crate::processor::{Envelope, FrameProcessor, Link};
pub trait WakeClock: Send {
fn now(&self) -> Duration;
}
pub struct SystemWakeClock {
base: std::time::Instant,
}
impl Default for SystemWakeClock {
fn default() -> Self {
Self {
base: std::time::Instant::now(),
}
}
}
impl WakeClock for SystemWakeClock {
fn now(&self) -> Duration {
self.base.elapsed()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum WakeState {
Idle,
Awake,
}
struct ParticipantState {
state: WakeState,
wake_timer: Duration,
accumulator: String,
}
impl ParticipantState {
fn new() -> Self {
Self {
state: WakeState::Idle,
wake_timer: Duration::ZERO,
accumulator: String::new(),
}
}
}
pub struct WakeCheckFilter {
name: &'static str,
wake_phrases: Vec<Vec<String>>,
keepalive: Duration,
clock: Box<dyn WakeClock>,
participants: HashMap<String, ParticipantState>,
}
impl WakeCheckFilter {
pub fn new(wake_phrases: Vec<String>, keepalive: Duration) -> Self {
Self::with_clock(
wake_phrases,
keepalive,
Box::new(SystemWakeClock::default()),
)
}
pub fn with_clock(
wake_phrases: Vec<String>,
keepalive: Duration,
clock: Box<dyn WakeClock>,
) -> Self {
let phrases = wake_phrases
.into_iter()
.map(|p| p.split_whitespace().map(|w| w.to_lowercase()).collect())
.collect();
Self {
name: "wake-check",
wake_phrases: phrases,
keepalive,
clock,
participants: HashMap::new(),
}
}
}
fn match_wake(wake_phrases: &[Vec<String>], haystack: &str) -> Option<usize> {
let lower = haystack.to_lowercase();
let words: Vec<(usize, &str)> = word_offsets(&lower);
for phrase in wake_phrases {
if phrase.is_empty() {
continue;
}
'start: for start in 0..words.len() {
if words.len() - start < phrase.len() {
break;
}
for (i, pw) in phrase.iter().enumerate() {
if words[start + i].1 != pw.as_str() {
continue 'start;
}
}
return Some(words[start].0);
}
}
None
}
fn word_offsets(s: &str) -> Vec<(usize, &str)> {
let mut out = Vec::new();
let mut idx = 0;
for word in s.split_whitespace() {
if let Some(rel) = s[idx..].find(word) {
let at = idx + rel;
out.push((at, word));
idx = at + word.len();
}
}
out
}
#[async_trait]
impl FrameProcessor for WakeCheckFilter {
fn name(&self) -> &str {
self.name
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
let Frame::Transcription {
ref text,
ref user_id,
..
} = env.frame
else {
link.push(env.meta, env.frame, env.direction).await;
return Ok(());
};
let now = self.clock.now();
let keepalive = self.keepalive;
let user_key = user_id.to_string();
let wake_phrases = &self.wake_phrases;
let p = self
.participants
.entry(user_key)
.or_insert_with(ParticipantState::new);
if p.state == WakeState::Awake {
if now.saturating_sub(p.wake_timer) < keepalive {
p.wake_timer = now;
link.push(env.meta, env.frame, env.direction).await;
return Ok(());
}
p.state = WakeState::Idle;
}
if !p.accumulator.is_empty() && !p.accumulator.ends_with(char::is_whitespace) {
p.accumulator.push(' ');
}
p.accumulator.push_str(text);
if let Some(start) = match_wake(wake_phrases, &p.accumulator) {
p.state = WakeState::Awake;
p.wake_timer = now;
let gated = p.accumulator[start..].to_string();
p.accumulator.clear();
let (uid, lang, final_) = match env.frame {
Frame::Transcription {
user_id,
language,
final_,
..
} => (user_id, language, final_),
_ => unreachable!(),
};
link.push(
env.meta,
Frame::Transcription {
text: gated,
user_id: uid,
language: lang,
final_,
},
env.direction,
)
.await;
}
Ok(())
}
}
pub type WakeNotifier = Arc<dyn Fn() + Send + Sync>;
pub type WakePredicate = Arc<dyn Fn(&Frame) -> bool + Send + Sync>;
pub struct WakeNotifierProcessor {
name: &'static str,
kinds: Vec<FrameKind>,
predicate: WakePredicate,
notifier: WakeNotifier,
}
impl WakeNotifierProcessor {
pub fn new(kinds: Vec<FrameKind>, predicate: WakePredicate, notifier: WakeNotifier) -> Self {
Self {
name: "wake-notifier",
kinds,
predicate,
notifier,
}
}
}
#[async_trait]
impl FrameProcessor for WakeNotifierProcessor {
fn name(&self) -> &str {
self.name
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if self.kinds.contains(&env.frame.kind()) && (self.predicate)(&env.frame) {
(self.notifier)();
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::test_harness::drive;
use crate::processor::frame::Direction;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Clone, Default)]
struct FakeClock(Arc<AtomicU64>);
impl FakeClock {
fn advance(&self, ms: u64) {
self.0.fetch_add(ms, Ordering::SeqCst);
}
}
impl WakeClock for FakeClock {
fn now(&self) -> Duration {
Duration::from_millis(self.0.load(Ordering::SeqCst))
}
}
fn transcription(text: &str, user: &str) -> Frame {
Frame::Transcription {
text: text.into(),
user_id: Arc::from(user),
language: None,
final_: true,
}
}
#[tokio::test]
async fn gates_until_wake_phrase_then_opens() {
let filter = WakeCheckFilter::new(vec!["hey assistant".into()], Duration::from_secs(3));
let out = drive(
Box::new(filter),
vec![
transcription("what time is it", "u1"), transcription("hey assistant set a timer", "u1"), transcription("for ten minutes", "u1"), ],
Direction::Downstream,
)
.await;
let texts: Vec<String> = out
.into_iter()
.filter_map(|f| match f {
Frame::Transcription { text, .. } => Some(text),
_ => None,
})
.collect();
assert_eq!(texts, vec!["hey assistant set a timer", "for ten minutes"]);
}
#[tokio::test]
async fn wake_only_fires_on_the_phrase() {
let filter = WakeCheckFilter::new(vec!["computer".into()], Duration::from_secs(3));
let out = drive(
Box::new(filter),
vec![
transcription("hello world", "u1"),
transcription("the computers are noisy", "u1"),
],
Direction::Downstream,
)
.await;
let count = out
.iter()
.filter(|f| matches!(f, Frame::Transcription { .. }))
.count();
assert_eq!(count, 0, "fired on a non-wake word: {out:?}");
}
#[tokio::test]
async fn keepalive_expires_and_regates() {
let clock = FakeClock::default();
let filter = WakeCheckFilter::with_clock(
vec!["wake".into()],
Duration::from_millis(1000),
Box::new(clock.clone()),
);
use crate::processor::runtime::{channel, NORMAL_CHAN_CAP};
use crate::processor::{Envelope, Link};
use std::sync::atomic::AtomicI64;
let (cap_tx, mut cap_rx) = channel(Arc::from("cap"), NORMAL_CHAN_CAP);
let link = Link {
next: Some(cap_tx.clone()),
prev: Some(cap_tx),
name: Arc::from("wake-check"),
clock: crate::processor::Clock::new(),
observer: None,
enable_metrics: false,
enable_usage_metrics: false,
ttfb_start: Arc::new(AtomicI64::new(0)),
processing_start: Arc::new(AtomicI64::new(0)),
};
let mut filter: Box<dyn FrameProcessor> = Box::new(filter);
let drain = |rx: &mut crate::processor::runtime::ProcessorRx| -> Vec<Frame> {
let mut v = Vec::new();
while let Ok(e) = rx.normal.try_recv() {
v.push(e.frame);
}
v
};
filter
.process_frame(
Envelope::new(transcription("wake up", "u"), Direction::Downstream),
&link,
)
.await
.unwrap();
assert_eq!(drain(&mut cap_rx).len(), 1);
clock.advance(500);
filter
.process_frame(
Envelope::new(transcription("still here", "u"), Direction::Downstream),
&link,
)
.await
.unwrap();
assert_eq!(drain(&mut cap_rx).len(), 1);
clock.advance(2000);
filter
.process_frame(
Envelope::new(transcription("ignored now", "u"), Direction::Downstream),
&link,
)
.await
.unwrap();
assert_eq!(
drain(&mut cap_rx).len(),
0,
"should re-gate after keepalive"
);
}
#[tokio::test]
async fn notifier_fires_on_matching_kind_and_predicate() {
let hits = Arc::new(AtomicU64::new(0));
let h2 = hits.clone();
let proc = WakeNotifierProcessor::new(
vec![FrameKind::Transcription],
Arc::new(|f| matches!(f, Frame::Transcription { text, .. } if text.contains("now"))),
Arc::new(move || {
h2.fetch_add(1, Ordering::SeqCst);
}),
);
let out = drive(
Box::new(proc),
vec![
transcription("later", "u"),
transcription("do it now", "u"),
Frame::Text("now".into()), ],
Direction::Downstream,
)
.await;
assert_eq!(hits.load(Ordering::SeqCst), 1);
assert_eq!(out.len(), 3, "notifier must forward every frame");
}
}