Skip to main content

substrate_memory/
lib.rs

1#![forbid(unsafe_code)]
2#![warn(missing_docs)]
3
4//! Two-tier agent memory: bounded ring buffer + persistent SQLite history.
5
6use std::collections::VecDeque;
7use std::sync::Mutex;
8
9use chrono::Utc;
10use store_sqlite::SqliteMemoryStore;
11use substrate_core::memory_port::{MemoryEntry, MemoryPort};
12use uuid::Uuid;
13
14/// Error type for in-memory memory adapters.
15#[derive(Debug, thiserror::Error)]
16pub enum MemoryError {
17    /// Underlying SQLite store failure.
18    #[error("store: {0}")]
19    Store(#[from] store_sqlite::StoreError),
20    /// Generic memory error.
21    #[error("{0}")]
22    Other(String),
23}
24
25/// Bounded ring buffer keeping the most recent `capacity` entries.
26pub struct RingMemory {
27    capacity: usize,
28    entries: Mutex<VecDeque<MemoryEntry>>,
29}
30
31impl RingMemory {
32    /// Create a ring buffer that retains at most `capacity` entries.
33    pub fn new(capacity: usize) -> Self {
34        Self {
35            capacity,
36            entries: Mutex::new(VecDeque::new()),
37        }
38    }
39
40    fn push_entry(&self, key: &str, content: &str) -> Uuid {
41        let id = Uuid::new_v4();
42        let entry = MemoryEntry {
43            id,
44            key: key.to_string(),
45            content: content.to_string(),
46            created_at: Utc::now().timestamp(),
47        };
48        let mut buf = self.entries.lock().unwrap();
49        if self.capacity > 0 && buf.len() >= self.capacity {
50            buf.pop_front();
51        }
52        buf.push_back(entry);
53        id
54    }
55}
56
57impl MemoryPort for RingMemory {
58    type Error = MemoryError;
59
60    fn append(&self, key: &str, content: &str) -> Result<Uuid, Self::Error> {
61        Ok(self.push_entry(key, content))
62    }
63
64    fn get(&self, key: &str) -> Result<Option<String>, Self::Error> {
65        let buf = self.entries.lock().unwrap();
66        Ok(buf
67            .iter()
68            .rev()
69            .find(|e| e.key == key)
70            .map(|e| e.content.clone()))
71    }
72
73    fn recent(&self, limit: usize) -> Result<Vec<MemoryEntry>, Self::Error> {
74        let buf = self.entries.lock().unwrap();
75        Ok(buf.iter().rev().take(limit).cloned().collect())
76    }
77
78    fn history(&self) -> Result<Vec<MemoryEntry>, Self::Error> {
79        self.recent(usize::MAX)
80    }
81}
82
83/// Composes a hot [`RingMemory`] tier with a cold [`SqliteMemoryStore`] tier.
84pub struct TwoTierMemory {
85    ring: RingMemory,
86    persistent: SqliteMemoryStore,
87}
88
89impl TwoTierMemory {
90    /// Create a two-tier store with ring capacity `ring_capacity`.
91    pub fn in_memory(ring_capacity: usize) -> Result<Self, MemoryError> {
92        Ok(Self {
93            ring: RingMemory::new(ring_capacity),
94            persistent: SqliteMemoryStore::open_in_memory()?,
95        })
96    }
97}
98
99impl MemoryPort for TwoTierMemory {
100    type Error = MemoryError;
101
102    fn append(&self, key: &str, content: &str) -> Result<Uuid, Self::Error> {
103        self.ring.append(key, content)?;
104        self.persistent
105            .append(key, content)
106            .map_err(MemoryError::from)
107    }
108
109    fn get(&self, key: &str) -> Result<Option<String>, Self::Error> {
110        if let Some(v) = self.ring.get(key)? {
111            return Ok(Some(v));
112        }
113        self.persistent.get(key).map_err(MemoryError::from)
114    }
115
116    fn recent(&self, limit: usize) -> Result<Vec<MemoryEntry>, Self::Error> {
117        self.ring.recent(limit)
118    }
119
120    fn history(&self) -> Result<Vec<MemoryEntry>, Self::Error> {
121        self.persistent.history().map_err(MemoryError::from)
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn ring_evicts_oldest_at_capacity() {
131        let ring = RingMemory::new(2);
132        ring.append("a", "1").unwrap();
133        ring.append("b", "2").unwrap();
134        ring.append("c", "3").unwrap();
135        let recent = ring.recent(10).unwrap();
136        assert_eq!(recent.len(), 2);
137        assert_eq!(recent[0].content, "3");
138        assert_eq!(recent[1].content, "2");
139        assert!(recent.iter().all(|e| e.content != "1"));
140    }
141
142    #[test]
143    fn persistent_round_trip() {
144        let store = SqliteMemoryStore::open_in_memory().unwrap();
145        store.append("topic", "hello").unwrap();
146        assert_eq!(store.get("topic").unwrap(), Some("hello".into()));
147        let hist = store.history().unwrap();
148        assert_eq!(hist.len(), 1);
149        assert_eq!(hist[0].content, "hello");
150    }
151
152    #[test]
153    fn two_tier_compose() {
154        let mem = TwoTierMemory::in_memory(2).unwrap();
155        mem.append("k", "v1").unwrap();
156        mem.append("k", "v2").unwrap();
157        mem.append("k", "v3").unwrap();
158        assert_eq!(mem.get("k").unwrap(), Some("v3".into()));
159        assert_eq!(mem.recent(10).unwrap().len(), 2);
160        assert_eq!(mem.history().unwrap().len(), 3);
161    }
162}