use std::collections::HashMap;
use tokio::sync::mpsc;
use terraphim_types::capability::ProcessId;
use crate::{AgentSpawner, SpawnerError};
#[derive(Debug)]
pub struct MentionRouter {
agents: HashMap<String, ProcessId>,
mention_receiver: mpsc::Receiver<MentionEvent>,
mention_sender: mpsc::Sender<MentionEvent>,
}
#[derive(Debug, Clone)]
pub struct MentionEvent {
pub from: ProcessId,
pub target: String,
pub message: String,
}
impl MentionRouter {
pub fn new() -> Self {
let (mention_sender, mention_receiver) = mpsc::channel(100);
Self {
agents: HashMap::new(),
mention_receiver,
mention_sender,
}
}
pub fn register_agent(&mut self, agent_id: String, process_id: ProcessId) {
self.agents.insert(agent_id, process_id);
}
pub fn sender(&self) -> mpsc::Sender<MentionEvent> {
self.mention_sender.clone()
}
pub async fn route_mentions(&mut self, _spawner: &AgentSpawner) -> Result<(), SpawnerError> {
while let Some(event) = self.mention_receiver.recv().await {
tracing::info!(
from = %event.from,
target = event.target.as_str(),
message = event.message.as_str(),
"Routing mention"
);
if let Some(&target_pid) = self.agents.get(&event.target) {
tracing::debug!(
target_agent = event.target.as_str(),
target_pid = %target_pid,
"Target agent found"
);
} else {
tracing::warn!(
target_agent = event.target.as_str(),
"Target agent not found, message dropped"
);
}
}
Ok(())
}
}
impl Default for MentionRouter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_mention_router() {
let mut router = MentionRouter::new();
let sender = router.sender();
router.register_agent("kimiko".to_string(), ProcessId::new());
let event = MentionEvent {
from: ProcessId::new(),
target: "kimiko".to_string(),
message: "Hello!".to_string(),
};
sender.send(event).await.unwrap();
let received = router.mention_receiver.recv().await;
assert!(received.is_some());
let received = received.unwrap();
assert_eq!(received.target, "kimiko");
}
}