Skip to main content

daimon_core/
memory.rs

1//! Conversation memory trait. Implement [`Memory`] for custom backends;
2//! built-in implementations live in the `daimon` facade crate.
3
4use std::future::Future;
5use std::pin::Pin;
6
7use crate::error::Result;
8use crate::types::Message;
9
10/// Trait for conversation memory backends. Stores and retrieves messages for agent context.
11pub trait Memory: Send + Sync {
12    /// Appends a message to the history. Order is preserved.
13    ///
14    /// Takes a borrow so callers that keep the message (the agent runner
15    /// appends it to its working log after persisting) don't clone it per
16    /// iteration. Backends that store owned messages clone internally;
17    /// serializing backends (Redis, SQLite) never need ownership at all.
18    fn add_message(&self, message: &Message) -> impl Future<Output = Result<()>> + Send;
19
20    /// Returns all stored messages in order. Used to build context for the model.
21    fn get_messages(&self) -> impl Future<Output = Result<Vec<Message>>> + Send;
22
23    /// Removes all messages. Use when starting a new conversation.
24    fn clear(&self) -> impl Future<Output = Result<()>> + Send;
25}
26
27/// Object-safe wrapper for the `Memory` trait, enabling dynamic dispatch via `Arc<dyn ErasedMemory>`.
28pub trait ErasedMemory: Send + Sync {
29    fn add_message_erased<'a>(
30        &'a self,
31        message: &'a Message,
32    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
33
34    fn get_messages_erased<'a>(
35        &'a self,
36    ) -> Pin<Box<dyn Future<Output = Result<Vec<Message>>> + Send + 'a>>;
37
38    fn clear_erased<'a>(&'a self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
39}
40
41impl<T: Memory> ErasedMemory for T {
42    fn add_message_erased<'a>(
43        &'a self,
44        message: &'a Message,
45    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
46        Box::pin(self.add_message(message))
47    }
48
49    fn get_messages_erased<'a>(
50        &'a self,
51    ) -> Pin<Box<dyn Future<Output = Result<Vec<Message>>> + Send + 'a>> {
52        Box::pin(self.get_messages())
53    }
54
55    fn clear_erased<'a>(&'a self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
56        Box::pin(self.clear())
57    }
58}
59
60/// Shared ownership of memory via `Arc<dyn ErasedMemory>`. Used by the agent.
61pub type SharedMemory = std::sync::Arc<dyn ErasedMemory>;
62
63#[cfg(test)]
64mod tests {
65    use super::{Memory, SharedMemory};
66    use crate::{Message, Result, Role};
67    use std::sync::{Arc, Mutex};
68
69    /// A provider-crate-style Memory impl using only daimon_core items.
70    struct VecMemory(Mutex<Vec<Message>>);
71
72    impl Memory for VecMemory {
73        async fn add_message(&self, message: &Message) -> Result<()> {
74            self.0.lock().unwrap().push(message.clone());
75            Ok(())
76        }
77
78        async fn get_messages(&self) -> Result<Vec<Message>> {
79            Ok(self.0.lock().unwrap().clone())
80        }
81
82        async fn clear(&self) -> Result<()> {
83            self.0.lock().unwrap().clear();
84            Ok(())
85        }
86    }
87
88    #[tokio::test]
89    async fn memory_is_implementable_from_core_alone() {
90        let mem = VecMemory(Mutex::new(Vec::new()));
91        mem.add_message(&Message::user("hi")).await.unwrap();
92        assert_eq!(mem.get_messages().await.unwrap().len(), 1);
93        assert_eq!(mem.get_messages().await.unwrap()[0].role, Role::User);
94        mem.clear().await.unwrap();
95        assert!(mem.get_messages().await.unwrap().is_empty());
96
97        let shared: SharedMemory = Arc::new(VecMemory(Mutex::new(Vec::new())));
98        shared
99            .add_message_erased(&Message::user("x"))
100            .await
101            .unwrap();
102        assert_eq!(shared.get_messages_erased().await.unwrap().len(), 1);
103    }
104}