johnandersen777.bsky.social/asyncai

Switch repository

Loading…

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<EventType>),
    All(Vec<EventType>),
}

impl EventSpec {
    pub fn is_satisfied(&self, seen: &HashSet<EventType>) -> 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<dyn FnOnce(&DispatchState) -> Option<AgiAction> + Send + 'static>;

pub struct WaitingRule {
    pub spec: EventSpec,
    make_action: ActionFactory,
}

impl WaitingRule {
    pub fn new<F>(spec: EventSpec, make_action: F) -> Self
    where
        F: FnOnce(&DispatchState) -> Option<AgiAction> + Send + 'static,
    {
        Self {
            spec,
            make_action: Box::new(make_action),
        }
    }

    pub fn one<F>(event: EventType, make_action: F) -> Self
    where
        F: FnOnce(&DispatchState) -> Option<AgiAction> + Send + 'static,
    {
        Self::new(EventSpec::One(event), make_action)
    }

    pub fn any<F>(events: Vec<EventType>, make_action: F) -> Self
    where
        F: FnOnce(&DispatchState) -> Option<AgiAction> + Send + 'static,
    {
        Self::new(EventSpec::Any(events), make_action)
    }

    pub fn all<F>(events: Vec<EventType>, make_action: F) -> Self
    where
        F: FnOnce(&DispatchState) -> Option<AgiAction> + Send + 'static,
    {
        Self::new(EventSpec::All(events), make_action)
    }
}

#[derive(Default)]
pub struct WaitingEngine {
    rules: VecDeque<WaitingRule>,
    seen: HashSet<EventType>,
}

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<AgiAction> {
        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<AgiAction> {
        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());
    }
}