use std::collections::{HashMap, VecDeque}; #[derive(Debug, Clone)] pub struct ContextMessage { pub sender: String, pub content: String, pub timestamp: i64, } struct RoomContext { messages: VecDeque, max_size: usize, } impl RoomContext { fn new(max_size: usize) -> Self { Self { messages: VecDeque::with_capacity(max_size), max_size, } } fn add(&mut self, msg: ContextMessage) { if self.messages.len() >= self.max_size { self.messages.pop_front(); } self.messages.push_back(msg); } fn get(&self) -> Vec { self.messages.iter().cloned().collect() } } pub struct ConversationManager { rooms: HashMap, room_window: usize, dm_window: usize, max_rooms: usize, } impl ConversationManager { pub fn new(room_window: usize, dm_window: usize) -> Self { Self { rooms: HashMap::new(), room_window, dm_window, max_rooms: 500, // todo(sienna): make this configurable } } pub fn add_message(&mut self, room_id: &str, is_dm: bool, msg: ContextMessage) { let window = if is_dm { self.dm_window } else { self.room_window }; // Evict oldest room if at capacity if !self.rooms.contains_key(room_id) && self.rooms.len() >= self.max_rooms { // Remove the room with the oldest latest message let oldest = self .rooms .iter() .min_by_key(|(_, ctx)| ctx.messages.back().map(|m| m.timestamp).unwrap_or(0)) .map(|(k, _)| k.to_owned()); if let Some(key) = oldest { self.rooms.remove(&key); } } let ctx = self .rooms .entry(room_id.to_owned()) .or_insert_with(|| RoomContext::new(window)); ctx.add(msg); } pub fn get_context(&self, room_id: &str) -> Vec { self.rooms .get(room_id) .map(|ctx| ctx.get()) .unwrap_or_default() } pub fn room_count(&self) -> usize { self.rooms.len() } } #[cfg(test)] mod tests { use super::*; fn msg(sender: &str, content: &str, ts: i64) -> ContextMessage { ContextMessage { sender: sender.to_string(), content: content.to_string(), timestamp: ts, } } #[test] fn test_add_and_get_messages() { let mut cm = ConversationManager::new(5, 10); cm.add_message("!room1:x", false, msg("alice", "hello", 1)); cm.add_message("!room1:x", false, msg("bob", "hi", 2)); let ctx = cm.get_context("!room1:x"); assert_eq!(ctx.len(), 2); assert_eq!(ctx[0].sender, "alice"); assert_eq!(ctx[1].sender, "bob"); } #[test] fn test_empty_room_returns_empty() { let cm = ConversationManager::new(5, 10); let ctx = cm.get_context("!nonexistent:x"); assert!(ctx.is_empty()); } #[test] fn test_sliding_window_group_room() { let mut cm = ConversationManager::new(3, 10); for i in 0..5 { cm.add_message("!room:x", false, msg("user", &format!("msg{i}"), i)); } let ctx = cm.get_context("!room:x"); assert_eq!(ctx.len(), 3); // Should keep the last 3 assert_eq!(ctx[0].content, "msg2"); assert_eq!(ctx[1].content, "msg3"); assert_eq!(ctx[2].content, "msg4"); } #[test] fn test_sliding_window_dm_room() { let mut cm = ConversationManager::new(3, 5); for i in 0..7 { cm.add_message("!dm:x", true, msg("user", &format!("dm{i}"), i)); } let ctx = cm.get_context("!dm:x"); assert_eq!(ctx.len(), 5); assert_eq!(ctx[0].content, "dm2"); assert_eq!(ctx[4].content, "dm6"); } #[test] fn test_multiple_rooms_independent() { let mut cm = ConversationManager::new(5, 10); cm.add_message("!a:x", false, msg("alice", "in room a", 1)); cm.add_message("!b:x", false, msg("bob", "in room b", 2)); assert_eq!(cm.get_context("!a:x").len(), 1); assert_eq!(cm.get_context("!b:x").len(), 1); assert_eq!(cm.get_context("!a:x")[0].content, "in room a"); assert_eq!(cm.get_context("!b:x")[0].content, "in room b"); } #[test] fn test_lru_eviction_at_max_rooms() { // Create a manager with max_rooms = 500 (default), but we'll use a small one let mut cm = ConversationManager::new(5, 10); cm.max_rooms = 3; // Add 3 rooms cm.add_message("!room1:x", false, msg("a", "r1", 100)); cm.add_message("!room2:x", false, msg("b", "r2", 200)); cm.add_message("!room3:x", false, msg("c", "r3", 300)); assert_eq!(cm.room_count(), 3); // Adding a 4th room should evict the one with oldest latest message (room1, ts=100) cm.add_message("!room4:x", false, msg("d", "r4", 400)); assert_eq!(cm.room_count(), 3); assert!(cm.get_context("!room1:x").is_empty()); // evicted assert_eq!(cm.get_context("!room2:x").len(), 1); assert_eq!(cm.get_context("!room3:x").len(), 1); assert_eq!(cm.get_context("!room4:x").len(), 1); } #[test] fn test_existing_room_not_evicted() { let mut cm = ConversationManager::new(5, 10); cm.max_rooms = 2; cm.add_message("!room1:x", false, msg("a", "r1", 100)); cm.add_message("!room2:x", false, msg("b", "r2", 200)); // Adding to existing room should NOT trigger eviction cm.add_message("!room1:x", false, msg("a", "r1 again", 300)); assert_eq!(cm.room_count(), 2); assert_eq!(cm.get_context("!room1:x").len(), 2); } #[test] fn test_message_ordering_preserved() { let mut cm = ConversationManager::new(10, 10); cm.add_message("!r:x", false, msg("a", "first", 1)); cm.add_message("!r:x", false, msg("b", "second", 2)); cm.add_message("!r:x", false, msg("c", "third", 3)); let ctx = cm.get_context("!r:x"); assert_eq!(ctx[0].timestamp, 1); assert_eq!(ctx[1].timestamp, 2); assert_eq!(ctx[2].timestamp, 3); } }