telos_agent/
storage.rs

1//! Session storage for persisting and resuming agent conversations.
2//!
3//! The backend is pluggable via the [`Storage`] trait. Two built-in
4//! implementations:
5//! - [`JsonlStorage`] — one JSON line per message, on disk, under `<dir>/<session_id>.jsonl`.
6//! - [`NoopStorage`] — black-hole; useful for tests and ephemeral sessions.
7
8use async_trait::async_trait;
9use std::collections::HashMap;
10use std::path::PathBuf;
11use tokio::fs::OpenOptions;
12use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
13
14use crate::error::AgentError;
15use crate::model::message::Message;
16use crate::tools::api::FileReadRecord;
17
18/// Session metadata persisted alongside messages so a resumed session can
19/// restore turn IDs, metrics, and file-read state.
20#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
21pub struct SessionMetadata {
22    pub next_turn_id: u64,
23    pub total_input_tokens: usize,
24    pub total_output_tokens: usize,
25    #[serde(default)]
26    pub total_prompt_cache_hit_tokens: usize,
27    #[serde(default)]
28    pub total_prompt_cache_miss_tokens: usize,
29    pub total_tool_calls: usize,
30    pub total_tool_errors: usize,
31    pub total_iterations: usize,
32    pub compaction_count: usize,
33    pub turn_count: usize,
34    pub retry_count: usize,
35    pub read_file_state: HashMap<PathBuf, FileReadRecord>,
36}
37
38/// Storage backend for persisting agent sessions.
39///
40/// [`JsonlStorage`] writes one JSON line per message; [`NoopStorage`] discards everything.
41#[async_trait]
42pub trait Storage: Send + Sync + std::fmt::Debug {
43    /// Overwrite the stored conversation with a full snapshot.
44    async fn save_snapshot(&self, session_id: &str, messages: &[Message])
45    -> Result<(), AgentError>;
46    /// Append messages to the existing log (does not truncate).
47    async fn append(&self, session_id: &str, messages: &[Message]) -> Result<(), AgentError>;
48    /// Load all messages for a session. Returns an empty vec when the session is unknown.
49    async fn load(&self, session_id: &str) -> Result<Vec<Message>, AgentError>;
50    /// Persist session metadata. The default is a no-op for backwards compatibility.
51    async fn save_metadata(
52        &self,
53        _session_id: &str,
54        _metadata: &SessionMetadata,
55    ) -> Result<(), AgentError> {
56        Ok(())
57    }
58    /// Load session metadata. The default returns `None` for backwards compatibility.
59    async fn load_metadata(
60        &self,
61        _session_id: &str,
62    ) -> Result<Option<SessionMetadata>, AgentError> {
63        Ok(None)
64    }
65    /// Delete all persisted data for a session. The default is a no-op.
66    async fn delete(&self, _session_id: &str) -> Result<(), AgentError> {
67        Ok(())
68    }
69}
70
71/// On-disk JSONL backend. Each message is serialised to one line; the file is
72/// named `<session_id>.jsonl` inside the configured directory.
73#[derive(Debug)]
74pub struct JsonlStorage {
75    dir: PathBuf,
76}
77
78impl JsonlStorage {
79    /// Create the storage directory eagerly and return a handle.
80    pub fn new(dir: impl Into<PathBuf>) -> Result<Self, AgentError> {
81        let dir = dir.into();
82        std::fs::create_dir_all(&dir)
83            .map_err(|e| AgentError::Config(format!("failed to create storage directory: {e}")))?;
84        Ok(Self { dir })
85    }
86
87    /// Path on disk for the given session ID.
88    fn path(&self, session_id: &str) -> Result<PathBuf, AgentError> {
89        validate_session_id(session_id)?;
90        Ok(self.dir.join(format!("{session_id}.jsonl")))
91    }
92
93    /// Path on disk for the session metadata file.
94    fn metadata_path(&self, session_id: &str) -> Result<PathBuf, AgentError> {
95        validate_session_id(session_id)?;
96        Ok(self.dir.join(format!("{session_id}.metadata.json")))
97    }
98}
99
100/// Reject session IDs that could escape the storage directory.
101fn validate_session_id(session_id: &str) -> Result<(), AgentError> {
102    if session_id.is_empty() {
103        return Err(AgentError::Config("session_id cannot be empty".into()));
104    }
105    if !session_id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
106        return Err(AgentError::Config(format!(
107            "session_id contains invalid characters: {session_id}"
108        )));
109    }
110    Ok(())
111}
112
113#[async_trait]
114impl Storage for JsonlStorage {
115    async fn save_snapshot(
116        &self,
117        session_id: &str,
118        messages: &[Message],
119    ) -> Result<(), AgentError> {
120        tokio::fs::create_dir_all(&self.dir)
121            .await
122            .map_err(|e| AgentError::Config(format!("failed to create storage directory: {e}")))?;
123        let path = self.path(session_id)?;
124        // Write to a temporary file first, then atomically rename to the target
125        // on the same filesystem. A crash mid-write leaves the original intact.
126        let tmp_path = path.with_extension("jsonl.tmp");
127        let file = OpenOptions::new()
128            .create(true)
129            .truncate(true)
130            .write(true)
131            .open(&tmp_path)
132            .await
133            .map_err(|e| AgentError::Config(format!("storage open failed: {e}")))?;
134
135        let mut writer = BufWriter::new(file);
136        for msg in messages {
137            let line = serde_json::to_string(msg)
138                .map_err(|e| AgentError::Config(format!("serialize failed: {e}")))?;
139            writer
140                .write_all(line.as_bytes())
141                .await
142                .map_err(|e| AgentError::Config(format!("storage write failed: {e}")))?;
143            writer
144                .write_all(b"\n")
145                .await
146                .map_err(|e| AgentError::Config(format!("storage write failed: {e}")))?;
147        }
148        writer
149            .flush()
150            .await
151            .map_err(|e| AgentError::Config(format!("storage flush failed: {e}")))?;
152        // Atomically rename temp -> target (same-filesystem guarantee on Linux).
153        tokio::fs::rename(&tmp_path, &path)
154            .await
155            .map_err(|e| AgentError::Config(format!("storage rename failed: {e}")))?;
156        Ok(())
157    }
158
159    async fn append(&self, session_id: &str, messages: &[Message]) -> Result<(), AgentError> {
160        tokio::fs::create_dir_all(&self.dir)
161            .await
162            .map_err(|e| AgentError::Config(format!("failed to create storage directory: {e}")))?;
163        let path = self.path(session_id)?;
164        let file = OpenOptions::new()
165            .create(true)
166            .append(true)
167            .open(&path)
168            .await
169            .map_err(|e| AgentError::Config(format!("storage open failed: {e}")))?;
170
171        let mut writer = BufWriter::new(file);
172        for msg in messages {
173            let line = serde_json::to_string(msg)
174                .map_err(|e| AgentError::Config(format!("serialize failed: {e}")))?;
175            writer
176                .write_all(line.as_bytes())
177                .await
178                .map_err(|e| AgentError::Config(format!("storage write failed: {e}")))?;
179            writer
180                .write_all(b"\n")
181                .await
182                .map_err(|e| AgentError::Config(format!("storage write failed: {e}")))?;
183        }
184        writer
185            .flush()
186            .await
187            .map_err(|e| AgentError::Config(format!("storage flush failed: {e}")))?;
188        Ok(())
189    }
190
191    async fn load(&self, session_id: &str) -> Result<Vec<Message>, AgentError> {
192        let path = self.path(session_id)?;
193        if !path.exists() {
194            // Unknown session — treat as empty rather than an error so resume() can fall back.
195            return Ok(Vec::new());
196        }
197
198        let file = tokio::fs::File::open(&path)
199            .await
200            .map_err(|e| AgentError::Config(format!("storage open failed: {e}")))?;
201        let reader = BufReader::new(file);
202        let mut lines = reader.lines();
203
204        let mut messages = Vec::new();
205        while let Some(line) = lines
206            .next_line()
207            .await
208            .map_err(|e| AgentError::Config(format!("storage read failed: {e}")))?
209        {
210            if line.trim().is_empty() {
211                continue;
212            }
213            let msg: Message = serde_json::from_str(&line)
214                .map_err(|e| AgentError::Config(format!("deserialize failed: {e}")))?;
215            messages.push(msg);
216        }
217
218        Ok(messages)
219    }
220
221    async fn save_metadata(
222        &self,
223        session_id: &str,
224        metadata: &SessionMetadata,
225    ) -> Result<(), AgentError> {
226        tokio::fs::create_dir_all(&self.dir)
227            .await
228            .map_err(|e| AgentError::Config(format!("failed to create storage directory: {e}")))?;
229        let path = self.metadata_path(session_id)?;
230        let tmp_path = path.with_extension("metadata.json.tmp");
231        let bytes = serde_json::to_vec_pretty(metadata)
232            .map_err(|e| AgentError::Config(format!("serialize metadata failed: {e}")))?;
233        tokio::fs::write(&tmp_path, bytes)
234            .await
235            .map_err(|e| AgentError::Config(format!("storage metadata write failed: {e}")))?;
236        tokio::fs::rename(&tmp_path, &path)
237            .await
238            .map_err(|e| AgentError::Config(format!("storage metadata rename failed: {e}")))?;
239        Ok(())
240    }
241
242    async fn load_metadata(&self, session_id: &str) -> Result<Option<SessionMetadata>, AgentError> {
243        let path = self.metadata_path(session_id)?;
244        if !path.exists() {
245            return Ok(None);
246        }
247        let bytes = tokio::fs::read(&path)
248            .await
249            .map_err(|e| AgentError::Config(format!("storage metadata read failed: {e}")))?;
250        if bytes.is_empty() {
251            return Ok(None);
252        }
253        let metadata = serde_json::from_slice(&bytes)
254            .map_err(|e| AgentError::Config(format!("deserialize metadata failed: {e}")))?;
255        Ok(Some(metadata))
256    }
257
258    async fn delete(&self, session_id: &str) -> Result<(), AgentError> {
259        let jsonl_path = self.path(session_id)?;
260        if jsonl_path.exists() {
261            tokio::fs::remove_file(&jsonl_path)
262                .await
263                .map_err(|e| AgentError::Config(format!("storage delete failed: {e}")))?;
264        }
265        let metadata_path = self.metadata_path(session_id)?;
266        if metadata_path.exists() {
267            tokio::fs::remove_file(&metadata_path)
268                .await
269                .map_err(|e| AgentError::Config(format!("storage metadata delete failed: {e}")))?;
270        }
271        Ok(())
272    }
273}
274
275/// Storage backend that discards everything — useful for tests / ephemeral sessions.
276#[derive(Debug)]
277pub struct NoopStorage;
278
279#[async_trait]
280impl Storage for NoopStorage {
281    async fn save_snapshot(
282        &self,
283        _session_id: &str,
284        _messages: &[Message],
285    ) -> Result<(), AgentError> {
286        Ok(())
287    }
288
289    async fn append(&self, _session_id: &str, _messages: &[Message]) -> Result<(), AgentError> {
290        Ok(())
291    }
292
293    async fn load(&self, _session_id: &str) -> Result<Vec<Message>, AgentError> {
294        Ok(Vec::new())
295    }
296
297    async fn save_metadata(
298        &self,
299        _session_id: &str,
300        _metadata: &SessionMetadata,
301    ) -> Result<(), AgentError> {
302        Ok(())
303    }
304
305    async fn load_metadata(
306        &self,
307        _session_id: &str,
308    ) -> Result<Option<SessionMetadata>, AgentError> {
309        Ok(None)
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[tokio::test]
318    async fn jsonl_roundtrip_save_and_load() {
319        let dir = std::env::temp_dir().join("tiny_agent_test_storage_roundtrip");
320        let _ = std::fs::remove_dir_all(&dir);
321        let storage = JsonlStorage::new(&dir).unwrap();
322
323        let msgs =
324            vec![Message::system("system"), Message::user("hello"), Message::assistant("hi there")];
325
326        storage.save_snapshot("test-session", &msgs).await.unwrap();
327        let loaded = storage.load("test-session").await.unwrap();
328        assert_eq!(loaded.len(), 3);
329        assert_eq!(loaded[0].text_content(), "system");
330        assert_eq!(loaded[1].text_content(), "hello");
331        assert_eq!(loaded[2].text_content(), "hi there");
332
333        let _ = std::fs::remove_dir_all(&dir);
334    }
335
336    #[tokio::test]
337    async fn jsonl_load_unknown_session_returns_empty() {
338        let storage =
339            JsonlStorage::new(std::env::temp_dir().join("tiny_agent_test_storage_unknown"))
340                .unwrap();
341        let loaded = storage.load("nonexistent-session").await.unwrap();
342        assert!(loaded.is_empty());
343    }
344
345    #[tokio::test]
346    async fn jsonl_append_preserves_existing_messages() {
347        let dir = std::env::temp_dir().join("tiny_agent_test_storage_append");
348        let _ = std::fs::remove_dir_all(&dir);
349        let storage = JsonlStorage::new(&dir).unwrap();
350
351        storage.save_snapshot("s", &[Message::user("first")]).await.unwrap();
352        storage.append("s", &[Message::assistant("second")]).await.unwrap();
353        let loaded = storage.load("s").await.unwrap();
354        assert_eq!(loaded.len(), 2);
355        assert_eq!(loaded[0].text_content(), "first"); // snapshot was truncated, so snapshot wrote it
356        assert_eq!(loaded[1].text_content(), "second");
357
358        let _ = std::fs::remove_dir_all(&dir);
359    }
360
361    #[tokio::test]
362    async fn jsonl_snapshot_replaces_content() {
363        let dir = std::env::temp_dir().join("tiny_agent_test_storage_snapshot");
364        let _ = std::fs::remove_dir_all(&dir);
365        let storage = JsonlStorage::new(&dir).unwrap();
366
367        storage.save_snapshot("s", &[Message::user("old")]).await.unwrap();
368        storage.save_snapshot("s", &[Message::user("new")]).await.unwrap();
369        let loaded = storage.load("s").await.unwrap();
370        assert_eq!(loaded.len(), 1);
371        assert_eq!(loaded[0].text_content(), "new");
372
373        let _ = std::fs::remove_dir_all(&dir);
374    }
375
376    #[tokio::test]
377    async fn noop_storage_always_returns_empty() {
378        let storage = NoopStorage;
379        storage.save_snapshot("x", &[Message::user("hi")]).await.unwrap();
380        let loaded = storage.load("x").await.unwrap();
381        assert!(loaded.is_empty());
382        storage.append("x", &[Message::user("more")]).await.unwrap();
383        let loaded2 = storage.load("x").await.unwrap();
384        assert!(loaded2.is_empty());
385    }
386
387    #[tokio::test]
388    async fn jsonl_rejects_path_traversal_session_id() {
389        let dir = std::env::temp_dir().join("tiny_agent_test_storage_path_traversal");
390        let _ = std::fs::remove_dir_all(&dir);
391        let storage = JsonlStorage::new(&dir).unwrap();
392
393        let result = storage.save_snapshot("../../../etc/evil", &[Message::user("x")]).await;
394        assert!(matches!(result, Err(AgentError::Config(_))));
395
396        let result = storage.load("dir/sub").await;
397        assert!(matches!(result, Err(AgentError::Config(_))));
398
399        let _ = std::fs::remove_dir_all(&dir);
400    }
401
402    #[tokio::test]
403    async fn jsonl_metadata_roundtrip() {
404        let dir = std::env::temp_dir().join("tiny_agent_test_storage_metadata");
405        let _ = std::fs::remove_dir_all(&dir);
406        let storage = JsonlStorage::new(&dir).unwrap();
407
408        let mut read_file_state = HashMap::new();
409        read_file_state.insert(
410            PathBuf::from("src/lib.rs"),
411            FileReadRecord {
412                content: "fn main() {}".to_string(),
413                timestamp_ms: 1234,
414                is_partial_view: true,
415                offset: Some(0),
416                limit: Some(10),
417            },
418        );
419
420        let metadata = SessionMetadata {
421            next_turn_id: 42,
422            total_input_tokens: 100,
423            total_output_tokens: 50,
424            total_prompt_cache_hit_tokens: 10,
425            total_prompt_cache_miss_tokens: 90,
426            total_tool_calls: 5,
427            total_tool_errors: 1,
428            total_iterations: 20,
429            compaction_count: 2,
430            turn_count: 10,
431            retry_count: 3,
432            read_file_state,
433        };
434
435        storage.save_metadata("s", &metadata).await.unwrap();
436        let loaded = storage.load_metadata("s").await.unwrap().unwrap();
437        assert_eq!(loaded.next_turn_id, 42);
438        assert_eq!(loaded.total_input_tokens, 100);
439        assert_eq!(loaded.total_prompt_cache_hit_tokens, 10);
440        assert_eq!(loaded.total_prompt_cache_miss_tokens, 90);
441        assert_eq!(loaded.total_tool_calls, 5);
442        let key = PathBuf::from("src/lib.rs");
443        assert!(loaded.read_file_state.contains_key(&key));
444        assert_eq!(loaded.read_file_state[&key].content, "fn main() {}");
445
446        let _ = std::fs::remove_dir_all(&dir);
447    }
448
449    #[tokio::test]
450    async fn jsonl_metadata_load_unknown_returns_none() {
451        let storage = JsonlStorage::new(
452            std::env::temp_dir().join("tiny_agent_test_storage_metadata_unknown"),
453        )
454        .unwrap();
455        let loaded = storage.load_metadata("nonexistent-session").await.unwrap();
456        assert!(loaded.is_none());
457    }
458}
telos_agent/storage.rs

telos_agent/
storage.rs
