microagents_storage/
memory.rs

1use microagents_events::{AgentEventAny, SessionInitEvent, types::AgentEvent};
2use std::{collections::HashMap, sync::Arc};
3use tokio::sync::RwLock;
4
5use crate::types::AgentStorage;
6
7/// In-memory implementation of [`AgentStorage`].
8///
9/// All data is lost when the process exits.
10#[derive(Debug)]
11pub struct InMemoryAgentStorage {
12    sessions: Arc<RwLock<HashMap<String, Vec<AgentEventAny>>>>,
13}
14
15impl Default for InMemoryAgentStorage {
16    fn default() -> Self {
17        Self {
18            sessions: Arc::new(RwLock::new(HashMap::new())),
19        }
20    }
21}
22
23#[async_trait::async_trait]
24impl AgentStorage for InMemoryAgentStorage {
25    async fn create_session(&self, event: SessionInitEvent) -> anyhow::Result<()> {
26        let mut sessions = self.sessions.write().await;
27        sessions.insert(
28            event.session_id.clone(),
29            vec![AgentEventAny::SessionInit(event)],
30        );
31        Ok(())
32    }
33
34    async fn update_session(&self, event: AgentEventAny) -> anyhow::Result<()> {
35        let session_id = &event.session_id();
36        let mut sessions = self.sessions.write().await;
37        let session = sessions.get_mut(session_id);
38        if let Some(s) = session {
39            s.push(event);
40            return Ok(());
41        }
42        Err(anyhow::anyhow!(
43            "Could not find {session_id} among the registered sessions"
44        ))
45    }
46
47    async fn get_session(&self, session_id: &str) -> anyhow::Result<Vec<AgentEventAny>> {
48        let sessions = self.sessions.read().await;
49        let session = sessions.get(session_id);
50        if let Some(s) = session {
51            let mut events = s.to_owned();
52            events.sort_by_key(|a| a.timestamp());
53            return Ok(events);
54        }
55        Err(anyhow::anyhow!(
56            "Could not find {session_id} among the registered sessions"
57        ))
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use chrono::Utc;
64    use microagents_events::{
65        AssistantResponseEvent, SessionStopEvent, Usage, UserPromptSubmitEvent,
66    };
67
68    use super::*;
69
70    #[tokio::test]
71    async fn test_default_init() {
72        let memory = InMemoryAgentStorage::default();
73        assert_eq!(memory.sessions.read().await.len(), 0);
74    }
75
76    #[tokio::test]
77    async fn test_create_session() {
78        let memory = InMemoryAgentStorage::default();
79        memory
80            .create_session(SessionInitEvent {
81                session_id: "1".to_string(),
82                model: "gpt-5.5".into(),
83                provider: "openai".into(),
84                system: "you are a helpful assistant".into(),
85                init_type: microagents_events::SessionInitType::Start,
86                timestamp: Utc::now(),
87            })
88            .await
89            .expect("Should be able to create a session");
90        let sessions = memory.sessions.read().await;
91        assert!(sessions.get("1").is_some_and(|v| {
92            v.len() == 1
93                && v.first()
94                    .is_some_and(|f| f.clone().to_jsonrpc().method == "session.init")
95        }));
96    }
97
98    #[tokio::test]
99    async fn test_create_update_get_session() {
100        let memory = InMemoryAgentStorage::default();
101        memory
102            .create_session(SessionInitEvent {
103                session_id: "1".to_string(),
104                model: "gpt-5.5".into(),
105                provider: "openai".into(),
106                system: "you are a helpful assistant".into(),
107                init_type: microagents_events::SessionInitType::Start,
108                timestamp: Utc::now(),
109            })
110            .await
111            .expect("Should be able to create a session");
112        memory
113            .update_session(AgentEventAny::UserPromptSubmit(UserPromptSubmitEvent {
114                prompt: "hello".to_string(),
115                session_id: "1".to_string(),
116                turn_id: "t1".to_string(),
117                timestamp: Utc::now(),
118            }))
119            .await
120            .expect("Should be able to update memory");
121        memory
122            .update_session(AgentEventAny::AssistantResponse(AssistantResponseEvent {
123                session_id: "1".to_string(),
124                turn_id: "t1".to_string(),
125                full_text: "hello".to_string(),
126                tool_calls: None,
127                timestamp: Utc::now(),
128            }))
129            .await
130            .expect("Should be able to update memory");
131        memory
132            .update_session(AgentEventAny::SessionStop(SessionStopEvent {
133                session_id: "1".to_string(),
134                result: Some("hello".to_string()),
135                error: None,
136                success: true,
137                timestamp: Utc::now(),
138                usage: Usage::default(),
139            }))
140            .await
141            .expect("Should be able to update memory");
142        let events = memory
143            .get_session("1")
144            .await
145            .expect("Should be able to get the session");
146        assert_eq!(events.len(), 4);
147        assert_eq!(events[0].to_jsonrpc().method, "session.init".to_string());
148        assert_eq!(
149            events[1].to_jsonrpc().method,
150            "user.prompt.submit".to_string()
151        );
152        assert_eq!(
153            events[2].to_jsonrpc().method,
154            "assistant.response".to_string()
155        );
156        assert_eq!(events[3].to_jsonrpc().method, "session.stop".to_string());
157    }
158}
microagents_storage/memory.rs

microagents_storage/
memory.rs
