use std::collections::{HashSet, VecDeque}; use serde::{Deserialize, Serialize}; use crate::actions::AgiAction; use crate::events::EventType; use crate::state::DispatchState; #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "kind", content = "events", rename_all = "snake_case")] pub enum EventSpec { One(EventType), Any(Vec), All(Vec), } impl EventSpec { pub fn is_satisfied(&self, seen: &HashSet) -> bool { match self { Self::One(event) => seen.contains(event), Self::Any(events) => events.iter().any(|event| seen.contains(event)), Self::All(events) => events.iter().all(|event| seen.contains(event)), } } } pub type ActionFactory = Box Option + Send + 'static>; pub struct WaitingRule { pub spec: EventSpec, make_action: ActionFactory, } impl WaitingRule { pub fn new(spec: EventSpec, make_action: F) -> Self where F: FnOnce(&DispatchState) -> Option + Send + 'static, { Self { spec, make_action: Box::new(make_action), } } pub fn one(event: EventType, make_action: F) -> Self where F: FnOnce(&DispatchState) -> Option + Send + 'static, { Self::new(EventSpec::One(event), make_action) } pub fn any(events: Vec, make_action: F) -> Self where F: FnOnce(&DispatchState) -> Option + Send + 'static, { Self::new(EventSpec::Any(events), make_action) } pub fn all(events: Vec, make_action: F) -> Self where F: FnOnce(&DispatchState) -> Option + Send + 'static, { Self::new(EventSpec::All(events), make_action) } } #[derive(Default)] pub struct WaitingEngine { rules: VecDeque, seen: HashSet, } impl WaitingEngine { pub fn new() -> Self { Self::default() } pub fn record(&mut self, event: EventType) { self.seen.insert(event); } pub fn forget(&mut self, event: EventType) { self.seen.remove(&event); } pub fn push(&mut self, rule: WaitingRule) { self.rules.push_back(rule); } pub fn len(&self) -> usize { self.rules.len() } pub fn is_empty(&self) -> bool { self.rules.is_empty() } /// Fire every rule whose condition currently holds. Fired rules are /// removed; the rest stay queued in their original order. pub fn drain_ready(&mut self, state: &DispatchState) -> Vec { let mut actions = Vec::new(); let mut pending = VecDeque::new(); while let Some(rule) = self.rules.pop_front() { if rule.spec.is_satisfied(&self.seen) { if let Some(action) = (rule.make_action)(state) { actions.push(action); } } else { pending.push_back(rule); } } self.rules = pending; actions } } #[cfg(test)] mod tests { use super::*; use crate::actions::{NewThreadAction, RunThreadAction}; fn new_thread_action(state: &DispatchState) -> Option { state .active_agent .clone() .map(|agent_id| AgiAction::NewThread(NewThreadAction { agent_id })) } #[test] fn one_event_matches_after_it_is_seen() { let mut engine = WaitingEngine::new(); engine.push(WaitingRule::one( EventType::NewAgentCreated, new_thread_action, )); engine.record(EventType::NewThreadCreated); assert!(engine .drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: None, }) .is_empty()); engine.record(EventType::NewAgentCreated); let actions = engine.drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: None, }); assert_eq!( actions, vec![AgiAction::NewThread(NewThreadAction { agent_id: "a1".to_string(), })] ); assert!(engine.is_empty()); } #[test] fn any_and_all_specs_work() { let mut any_engine = WaitingEngine::new(); any_engine.push(WaitingRule::any( vec![ EventType::NewAgentCreated, EventType::ExistingAgentRetrieved, ], new_thread_action, )); any_engine.record(EventType::ExistingAgentRetrieved); assert_eq!( any_engine .drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: None, }) .len(), 1 ); let mut all_engine = WaitingEngine::new(); all_engine.push(WaitingRule::all( vec![EventType::NewThreadCreated, EventType::ThreadMessageAdded], |state| { Some(AgiAction::RunThread(RunThreadAction { agent_id: state.active_agent.clone()?, thread_id: state.active_thread.clone()?, })) }, )); all_engine.record(EventType::NewThreadCreated); assert!(all_engine .drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: Some("t1".to_string()), }) .is_empty()); all_engine.record(EventType::ThreadMessageAdded); assert_eq!( all_engine .drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: Some("t1".to_string()), }) .len(), 1 ); } #[test] fn forget_blocks_stale_matches() { let mut engine = WaitingEngine::new(); engine.record(EventType::ThreadMessageAdded); engine.push(WaitingRule::one( EventType::ThreadMessageAdded, new_thread_action, )); engine.forget(EventType::ThreadMessageAdded); assert!(engine .drain_ready(&DispatchState { active_agent: Some("a1".to_string()), active_thread: None, }) .is_empty()); assert_eq!(engine.len(), 1); } #[test] fn rules_fire_in_insertion_order() { let mut engine = WaitingEngine::new(); engine.record(EventType::NewAgentCreated); let state = DispatchState { active_agent: Some("a1".to_string()), active_thread: None, }; engine.push(WaitingRule::one(EventType::NewAgentCreated, |_| { Some(AgiAction::NewThread(NewThreadAction { agent_id: "first".to_string(), })) })); engine.push(WaitingRule::one(EventType::NewAgentCreated, |_| { Some(AgiAction::NewThread(NewThreadAction { agent_id: "second".to_string(), })) })); let actions = engine.drain_ready(&state); assert_eq!( actions, vec![ AgiAction::NewThread(NewThreadAction { agent_id: "first".to_string(), }), AgiAction::NewThread(NewThreadAction { agent_id: "second".to_string(), }), ] ); } #[test] fn factory_returning_none_consumes_the_rule() { let mut engine = WaitingEngine::new(); engine.record(EventType::NewAgentCreated); engine.push(WaitingRule::one(EventType::NewAgentCreated, |_| None)); assert!(engine.drain_ready(&DispatchState::default()).is_empty()); assert!(engine.is_empty()); } }