rskit_agent/memory/
mod.rs

1//! Conversation memory — store and retrieve message history per session.
2//!
3//! The [`Memory`] trait defines async operations for loading, saving, and
4//! appending messages.  [`InMemoryStore`] keeps everything in a
5//! `parking_lot::RwLock<HashMap>`, while [`SlidingWindowMemory`] wraps any
6//! `Memory` and trims to a maximum message count.
7
8mod sliding_window;
9mod store;
10
11pub use sliding_window::SlidingWindowMemory;
12pub use store::InMemoryStore;
13
14use async_trait::async_trait;
15use rskit_errors::AppError;
16use rskit_llm::types::Message;
17
18// ── Memory trait ────────────────────────────────────────────────────────────
19
20/// Async conversation memory backend.
21#[async_trait]
22pub trait Memory: Send + Sync {
23    /// Load all messages for a session.
24    async fn load(&self, session_id: &str) -> Result<Vec<Message>, AppError>;
25
26    /// Replace the entire message list for a session.
27    async fn save(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError>;
28
29    /// Append messages to an existing session (creating it if needed).
30    async fn append(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError>;
31
32    /// Remove all messages for a session.
33    async fn clear(&self, session_id: &str) -> Result<(), AppError>;
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use rskit_llm::types;
40    use std::sync::Arc;
41
42    #[tokio::test]
43    async fn test_in_memory_store_save_load() {
44        let store = InMemoryStore::new();
45        let msgs = vec![types::user("hello"), types::assistant("hi")];
46        store.save("s1", &msgs).await.unwrap();
47
48        let loaded = store.load("s1").await.unwrap();
49        assert_eq!(loaded.len(), 2);
50    }
51
52    #[tokio::test]
53    async fn test_in_memory_store_append() {
54        let store = InMemoryStore::new();
55        store.save("s1", &[types::user("hello")]).await.unwrap();
56        store.append("s1", &[types::assistant("hi")]).await.unwrap();
57
58        let loaded = store.load("s1").await.unwrap();
59        assert_eq!(loaded.len(), 2);
60    }
61
62    #[tokio::test]
63    async fn test_in_memory_store_clear() {
64        let store = InMemoryStore::new();
65        store.save("s1", &[types::user("hello")]).await.unwrap();
66        store.clear("s1").await.unwrap();
67
68        let loaded = store.load("s1").await.unwrap();
69        assert!(loaded.is_empty());
70    }
71
72    #[tokio::test]
73    async fn test_in_memory_store_load_missing() {
74        let store = InMemoryStore::new();
75        let loaded = store.load("nonexistent").await.unwrap();
76        assert!(loaded.is_empty());
77    }
78
79    #[tokio::test]
80    async fn test_sliding_window_trims() {
81        let inner = Arc::new(InMemoryStore::new());
82        let sw = SlidingWindowMemory::new(inner.clone(), 2).unwrap();
83
84        let msgs = vec![
85            types::user("a"),
86            types::assistant("b"),
87            types::user("c"),
88            types::assistant("d"),
89        ];
90        sw.save("s1", &msgs).await.unwrap();
91
92        let loaded = sw.load("s1").await.unwrap();
93        assert_eq!(loaded.len(), 2);
94    }
95
96    #[tokio::test]
97    async fn test_sliding_window_preserves_system() {
98        let inner = Arc::new(InMemoryStore::new());
99        let sw = SlidingWindowMemory::new(inner.clone(), 2).unwrap();
100
101        let msgs = vec![
102            types::system("sys"),
103            types::user("a"),
104            types::assistant("b"),
105            types::user("c"),
106            types::assistant("d"),
107        ];
108        sw.save("s1", &msgs).await.unwrap();
109
110        let loaded = sw.load("s1").await.unwrap();
111        // system + last 2
112        assert_eq!(loaded.len(), 3);
113        assert!(matches!(&loaded[0], Message::System(_)));
114    }
115
116    #[tokio::test]
117    async fn test_sliding_window_append_trims() {
118        let inner = Arc::new(InMemoryStore::new());
119        let sw = SlidingWindowMemory::new(inner.clone(), 3).unwrap();
120
121        sw.save("s1", &[types::user("a"), types::assistant("b")])
122            .await
123            .unwrap();
124        sw.append("s1", &[types::user("c"), types::assistant("d")])
125            .await
126            .unwrap();
127
128        let loaded = sw.load("s1").await.unwrap();
129        assert_eq!(loaded.len(), 3);
130    }
131
132    #[test]
133    fn test_sliding_window_zero_max_messages() {
134        let inner = Arc::new(InMemoryStore::new());
135        let result = SlidingWindowMemory::new(inner, 0);
136        assert!(result.is_err());
137    }
138}
rskit_agent/memory/mod.rs

rskit_agent/memory/
mod.rs
