saorsa_agent/session/
storage.rs

1//! Session storage and serialization.
2
3use crate::SaorsaAgentError;
4use crate::session::{Message, SessionId, SessionMetadata, SessionNode};
5use std::fs;
6use std::path::PathBuf;
7
8/// Manages filesystem storage for sessions.
9#[derive(Clone)]
10pub struct SessionStorage {
11    base_path: PathBuf,
12}
13
14impl SessionStorage {
15    /// Create a new storage manager with the default base path.
16    pub fn new() -> Result<Self, SaorsaAgentError> {
17        let base_path = crate::session::path::sessions_dir()?;
18        Ok(Self { base_path })
19    }
20
21    /// Create a storage manager with a custom base path (for testing).
22    pub fn with_base_path(base_path: PathBuf) -> Self {
23        Self { base_path }
24    }
25
26    /// Get the base path for all sessions.
27    pub fn base_path(&self) -> &std::path::Path {
28        &self.base_path
29    }
30
31    /// Get the directory for a specific session.
32    fn session_dir(&self, session_id: &SessionId) -> PathBuf {
33        self.base_path.join(session_id.as_str())
34    }
35
36    /// Ensure the session directory exists.
37    fn ensure_session_dir(&self, session_id: &SessionId) -> Result<(), SaorsaAgentError> {
38        let dir = self.session_dir(session_id);
39        crate::session::path::ensure_dir(&dir)?;
40        crate::session::path::ensure_dir(&dir.join("messages"))?;
41        Ok(())
42    }
43
44    /// Write data to a file atomically (write to temp, then rename).
45    fn write_atomic(&self, path: &std::path::Path, data: &[u8]) -> Result<(), SaorsaAgentError> {
46        let temp_path = path.with_extension("tmp");
47        fs::write(&temp_path, data)
48            .map_err(|e| SaorsaAgentError::Session(format!("Failed to write temp file: {}", e)))?;
49        fs::rename(&temp_path, path)
50            .map_err(|e| SaorsaAgentError::Session(format!("Failed to rename temp file: {}", e)))?;
51        Ok(())
52    }
53
54    /// Save session metadata to manifest.json.
55    pub fn save_manifest(
56        &self,
57        session_id: &SessionId,
58        metadata: &SessionMetadata,
59    ) -> Result<(), SaorsaAgentError> {
60        self.ensure_session_dir(session_id)?;
61        let path = self.session_dir(session_id).join("manifest.json");
62        let json = serde_json::to_string_pretty(metadata).map_err(|e| {
63            SaorsaAgentError::Session(format!("Failed to serialize manifest: {}", e))
64        })?;
65        self.write_atomic(&path, json.as_bytes())?;
66        Ok(())
67    }
68
69    /// Load session metadata from manifest.json.
70    pub fn load_manifest(
71        &self,
72        session_id: &SessionId,
73    ) -> Result<SessionMetadata, SaorsaAgentError> {
74        let path = self.session_dir(session_id).join("manifest.json");
75        let json = fs::read_to_string(&path)
76            .map_err(|e| SaorsaAgentError::Session(format!("Failed to read manifest: {}", e)))?;
77        serde_json::from_str(&json).map_err(|e| {
78            SaorsaAgentError::Session(format!("Failed to deserialize manifest: {}", e))
79        })
80    }
81
82    /// Save session tree structure to tree.json.
83    pub fn save_tree(
84        &self,
85        session_id: &SessionId,
86        node: &SessionNode,
87    ) -> Result<(), SaorsaAgentError> {
88        self.ensure_session_dir(session_id)?;
89        let path = self.session_dir(session_id).join("tree.json");
90        let json = serde_json::to_string_pretty(node)
91            .map_err(|e| SaorsaAgentError::Session(format!("Failed to serialize tree: {}", e)))?;
92        self.write_atomic(&path, json.as_bytes())?;
93        Ok(())
94    }
95
96    /// Load session tree structure from tree.json.
97    pub fn load_tree(&self, session_id: &SessionId) -> Result<SessionNode, SaorsaAgentError> {
98        let path = self.session_dir(session_id).join("tree.json");
99        let json = fs::read_to_string(&path)
100            .map_err(|e| SaorsaAgentError::Session(format!("Failed to read tree: {}", e)))?;
101        serde_json::from_str(&json)
102            .map_err(|e| SaorsaAgentError::Session(format!("Failed to deserialize tree: {}", e)))
103    }
104
105    /// Save a message to messages/{index}-{type}.json.
106    pub fn save_message(
107        &self,
108        session_id: &SessionId,
109        index: usize,
110        message: &Message,
111    ) -> Result<(), SaorsaAgentError> {
112        self.ensure_session_dir(session_id)?;
113
114        let message_type = match message {
115            Message::User { .. } => "user",
116            Message::Assistant { .. } => "assistant",
117            Message::ToolCall { .. } => "tool_call",
118            Message::ToolResult { .. } => "tool_result",
119        };
120
121        let path = self
122            .session_dir(session_id)
123            .join("messages")
124            .join(format!("{}-{}.json", index, message_type));
125
126        let json = serde_json::to_string_pretty(message).map_err(|e| {
127            SaorsaAgentError::Session(format!("Failed to serialize message: {}", e))
128        })?;
129
130        self.write_atomic(&path, json.as_bytes())?;
131        Ok(())
132    }
133
134    /// Load all messages for a session, in order.
135    pub fn load_messages(&self, session_id: &SessionId) -> Result<Vec<Message>, SaorsaAgentError> {
136        let messages_dir = self.session_dir(session_id).join("messages");
137
138        if !messages_dir.exists() {
139            return Ok(Vec::new());
140        }
141
142        let mut entries: Vec<_> = fs::read_dir(&messages_dir)
143            .map_err(|e| {
144                SaorsaAgentError::Session(format!("Failed to read messages directory: {}", e))
145            })?
146            .collect::<Result<Vec<_>, _>>()
147            .map_err(|e| {
148                SaorsaAgentError::Session(format!("Failed to read directory entry: {}", e))
149            })?;
150
151        // Sort by index (extracted from filename)
152        entries.sort_by_key(|entry| {
153            entry
154                .file_name()
155                .to_string_lossy()
156                .split('-')
157                .next()
158                .and_then(|s| s.parse::<usize>().ok())
159                .unwrap_or(usize::MAX)
160        });
161
162        let mut messages = Vec::new();
163        for entry in entries {
164            let path = entry.path();
165            if path.extension().and_then(|s| s.to_str()) == Some("json") {
166                let json = fs::read_to_string(&path).map_err(|e| {
167                    SaorsaAgentError::Session(format!("Failed to read message file: {}", e))
168                })?;
169                let message: Message = serde_json::from_str(&json).map_err(|e| {
170                    SaorsaAgentError::Session(format!("Failed to deserialize message: {}", e))
171                })?;
172                messages.push(message);
173            }
174        }
175
176        Ok(messages)
177    }
178}
179
180impl Default for SessionStorage {
181    fn default() -> Self {
182        Self::new().unwrap_or_else(|_| {
183            Self::with_base_path(PathBuf::from("/tmp/saorsa-sessions-fallback"))
184        })
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use tempfile::TempDir;
192
193    fn test_storage() -> (TempDir, SessionStorage) {
194        let temp_dir = match TempDir::new() {
195            Ok(dir) => dir,
196            Err(_) => panic!("Failed to create temp dir for test"),
197        };
198        let storage = SessionStorage::with_base_path(temp_dir.path().to_path_buf());
199        (temp_dir, storage)
200    }
201
202    #[test]
203    fn test_ensure_session_dir_creates_directories() {
204        let (_temp, storage) = test_storage();
205        let id = SessionId::new();
206
207        assert!(storage.ensure_session_dir(&id).is_ok());
208
209        let session_dir = storage.session_dir(&id);
210        assert!(session_dir.exists());
211        assert!(session_dir.join("messages").exists());
212    }
213
214    #[test]
215    fn test_manifest_roundtrip() {
216        let (_temp, storage) = test_storage();
217        let id = SessionId::new();
218        let mut metadata = SessionMetadata::new();
219        metadata.title = Some("Test Session".to_string());
220        metadata.add_tag("rust".to_string());
221
222        assert!(storage.save_manifest(&id, &metadata).is_ok());
223        let loaded = storage.load_manifest(&id);
224        assert!(loaded.is_ok());
225        match loaded {
226            Ok(loaded_meta) => {
227                assert!(loaded_meta.title == metadata.title);
228                assert!(loaded_meta.tags == metadata.tags);
229            }
230            Err(_) => unreachable!(),
231        }
232    }
233
234    #[test]
235    fn test_tree_roundtrip() {
236        let (_temp, storage) = test_storage();
237        let id = SessionId::new();
238        let parent_id = SessionId::new();
239        let mut node = SessionNode::new_child(id, parent_id);
240        node.add_child(SessionId::new());
241
242        assert!(storage.save_tree(&id, &node).is_ok());
243        let loaded = storage.load_tree(&id);
244        assert!(loaded.is_ok());
245        match loaded {
246            Ok(loaded_node) => {
247                assert!(loaded_node.id == node.id);
248                assert!(loaded_node.parent_id == node.parent_id);
249                assert!(loaded_node.child_ids.len() == node.child_ids.len());
250            }
251            Err(_) => unreachable!(),
252        }
253    }
254
255    #[test]
256    fn test_message_serialization() {
257        let (_temp, storage) = test_storage();
258        let id = SessionId::new();
259
260        let msg1 = Message::user("Hello".to_string());
261        let msg2 = Message::assistant("Hi there".to_string());
262        let msg3 = Message::tool_call("bash".to_string(), serde_json::json!({"cmd": "ls"}));
263
264        assert!(storage.save_message(&id, 0, &msg1).is_ok());
265        assert!(storage.save_message(&id, 1, &msg2).is_ok());
266        assert!(storage.save_message(&id, 2, &msg3).is_ok());
267
268        let messages_dir = storage.session_dir(&id).join("messages");
269        assert!(messages_dir.join("0-user.json").exists());
270        assert!(messages_dir.join("1-assistant.json").exists());
271        assert!(messages_dir.join("2-tool_call.json").exists());
272    }
273
274    #[test]
275    fn test_load_messages_in_order() {
276        let (_temp, storage) = test_storage();
277        let id = SessionId::new();
278
279        let msg1 = Message::user("First".to_string());
280        let msg2 = Message::assistant("Second".to_string());
281        let msg3 = Message::user("Third".to_string());
282
283        assert!(storage.save_message(&id, 0, &msg1).is_ok());
284        assert!(storage.save_message(&id, 1, &msg2).is_ok());
285        assert!(storage.save_message(&id, 2, &msg3).is_ok());
286
287        let loaded = storage.load_messages(&id);
288        assert!(loaded.is_ok());
289        match loaded {
290            Ok(messages) => {
291                assert!(messages.len() == 3);
292
293                match &messages[0] {
294                    Message::User { content, .. } => assert!(content == "First"),
295                    _ => panic!("Expected User message"),
296                }
297
298                match &messages[1] {
299                    Message::Assistant { content, .. } => assert!(content == "Second"),
300                    _ => panic!("Expected Assistant message"),
301                }
302
303                match &messages[2] {
304                    Message::User { content, .. } => assert!(content == "Third"),
305                    _ => panic!("Expected User message"),
306                }
307            }
308            Err(_) => unreachable!(),
309        }
310    }
311
312    #[test]
313    fn test_load_messages_empty_session() {
314        let (_temp, storage) = test_storage();
315        let id = SessionId::new();
316
317        let messages = storage.load_messages(&id);
318        assert!(messages.is_ok());
319        match messages {
320            Ok(msgs) => assert!(msgs.is_empty()),
321            Err(_) => unreachable!(),
322        }
323    }
324
325    #[test]
326    fn test_atomic_write_creates_and_renames() {
327        let (_temp, storage) = test_storage();
328        let id = SessionId::new();
329
330        assert!(storage.ensure_session_dir(&id).is_ok());
331        let path = storage.session_dir(&id).join("test.json");
332
333        assert!(storage.write_atomic(&path, b"test data").is_ok());
334
335        assert!(path.exists());
336        assert!(!path.with_extension("tmp").exists());
337
338        let content = fs::read_to_string(&path);
339        assert!(content.is_ok());
340        match content {
341            Ok(c) => assert!(c == "test data"),
342            Err(_) => unreachable!(),
343        }
344    }
345}
saorsa_agent/session/storage.rs

saorsa_agent/session/
storage.rs
