llm_agent_runtime/
persistence.rs

1//! # Module: Persistence
2//!
3//! ## Responsibility
4//! Provides a pluggable persistence backend trait and a `FilePersistenceBackend`
5//! that stores serialized state as files on disk using `tokio::fs`.
6//!
7//! ## Guarantees
8//! - `PersistenceBackend` is object-safe via `async-trait`
9//! - `FilePersistenceBackend` stores each key as a file in a configurable directory
10//! - Non-panicking: all operations return `Result`
11//!
12//! ## Feature Gate
13//! This module is only compiled when the `persistence` feature is enabled.
14
15use crate::error::AgentRuntimeError;
16use crate::util::djb2;
17use async_trait::async_trait;
18use futures::future::try_join_all;
19use std::path::PathBuf;
20use std::sync::Arc;
21use uuid::Uuid;
22
23// ── PersistenceBackend ────────────────────────────────────────────────────────
24
25/// Pluggable storage backend for agent state (memory, graph, session checkpoints).
26///
27/// Implement this trait to persist agent state to any storage system
28/// (disk, Redis, S3, databases, etc.).
29#[async_trait]
30pub trait PersistenceBackend: Send + Sync {
31    /// Persist raw bytes under the given key.
32    async fn save(&self, key: &str, value: &[u8]) -> Result<(), AgentRuntimeError>;
33
34    /// Load raw bytes for the given key.
35    ///
36    /// Returns `None` if no data has been stored under this key.
37    async fn load(&self, key: &str) -> Result<Option<Vec<u8>>, AgentRuntimeError>;
38
39    /// Delete the entry for the given key. No-op if the key does not exist.
40    async fn delete(&self, key: &str) -> Result<(), AgentRuntimeError>;
41
42    /// Save multiple key-value pairs concurrently.
43    ///
44    /// The default implementation issues all `save` calls concurrently using
45    /// [`futures::future::try_join_all`] so backends that incur per-call I/O
46    /// latency benefit without any additional code.  Backends with their own
47    /// native batch API should override this method.
48    async fn batch_save(&self, items: &[(&str, &[u8])]) -> Result<(), AgentRuntimeError> {
49        try_join_all(items.iter().map(|(key, value)| self.save(key, value))).await?;
50        Ok(())
51    }
52
53    /// Load multiple keys concurrently. Returns a vec of `Option<Vec<u8>>` in the same order.
54    ///
55    /// The default implementation issues all `load` calls concurrently using
56    /// [`futures::future::try_join_all`].  Backends with a native multi-get
57    /// API should override this method.
58    async fn batch_load(&self, keys: &[&str]) -> Result<Vec<Option<Vec<u8>>>, AgentRuntimeError> {
59        try_join_all(keys.iter().map(|key| self.load(key))).await
60    }
61}
62
63// ── FilePersistenceBackend ────────────────────────────────────────────────────
64
65/// Persists data as files in a directory on disk.
66///
67/// Each key maps to a file named `<key>.bin` inside the base directory.
68/// The base directory must exist before calling any methods.
69///
70/// # Key Sanitization
71///
72/// Characters that are invalid in file names on common operating systems
73/// (`/`, `\`, `:`, `*`, `?`, `"`, `<`, `>`, `|`) are replaced with `_`
74/// before the path is constructed. This prevents path-traversal attacks
75/// and ensures portability across Linux, macOS, and Windows.
76///
77/// # Concurrency
78///
79/// `FilePersistenceBackend` is `Clone` and `Send + Sync`. Each clone
80/// shares the same `base_dir` via `Arc` and can be used from multiple
81/// async tasks simultaneously.
82///
83/// Writes are atomic at the OS level: data is written to a uniquely-named
84/// temporary file in the same directory and then renamed into place, so a
85/// reader never observes a half-written file even if the process crashes.
86/// Concurrent writes to the **same key** from multiple tasks are safe but
87/// are not serialized — the final file content is determined by whichever
88/// rename completes last.
89#[derive(Debug, Clone)]
90pub struct FilePersistenceBackend {
91    /// Absolute path to the directory where `<key>.bin` files are stored.
92    base_dir: Arc<PathBuf>,
93}
94
95impl FilePersistenceBackend {
96    /// Create a new backend that stores files in `base_dir`.
97    ///
98    /// The directory is not created automatically — it must already exist.
99    pub fn new(base_dir: impl Into<PathBuf>) -> Self {
100        Self {
101            base_dir: Arc::new(base_dir.into()),
102        }
103    }
104
105    /// Return the base directory that this backend stores files in.
106    pub fn base_dir(&self) -> &std::path::Path {
107        self.base_dir.as_ref()
108    }
109
110    /// Compute the file path for a given key.
111    ///
112    /// Uses a `<readable_prefix>-<djb2_hash>.bin` scheme to guarantee uniqueness
113    /// even when two keys sanitize to the same string (e.g. `"a/b"` and `"a_b"`).
114    /// The readable prefix aids manual inspection of the directory; the 16-digit
115    /// hex hash is the canonical disambiguator.
116    fn path_for(&self, key: &str) -> PathBuf {
117        let hash = djb2(key);
118        let sanitized = key.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "_");
119        self.base_dir
120            .join(format!("{sanitized}-{hash:016x}.bin"))
121    }
122}
123
124#[async_trait]
125impl PersistenceBackend for FilePersistenceBackend {
126    #[tracing::instrument(skip(self, value), fields(key))]
127    async fn save(&self, key: &str, value: &[u8]) -> Result<(), AgentRuntimeError> {
128        let path = self.path_for(key);
129        // Write to a unique temp file then atomically rename to the target path.
130        // This prevents a reader from observing a half-written file if the
131        // process crashes mid-write.
132        let tmp_path = path.with_extension(format!("tmp-{}", Uuid::new_v4().simple()));
133        tokio::fs::write(&tmp_path, value)
134            .await
135            .map_err(|e| AgentRuntimeError::Persistence(format!("write tmp {tmp_path:?}: {e}")))?;
136        tokio::fs::rename(&tmp_path, &path).await.map_err(|e| {
137            AgentRuntimeError::Persistence(format!(
138                "rename {tmp_path:?} -> {path:?}: {e}"
139            ))
140        })?;
141        Ok(())
142    }
143
144    #[tracing::instrument(skip(self), fields(key))]
145    async fn load(&self, key: &str) -> Result<Option<Vec<u8>>, AgentRuntimeError> {
146        let path = self.path_for(key);
147        match tokio::fs::read(&path).await {
148            Ok(bytes) => Ok(Some(bytes)),
149            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
150            Err(e) => Err(AgentRuntimeError::Persistence(format!(
151                "read {path:?}: {e}"
152            ))),
153        }
154    }
155
156    async fn delete(&self, key: &str) -> Result<(), AgentRuntimeError> {
157        let path = self.path_for(key);
158        match tokio::fs::remove_file(&path).await {
159            Ok(()) => Ok(()),
160            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
161            Err(e) => Err(AgentRuntimeError::Persistence(format!(
162                "delete {path:?}: {e}"
163            ))),
164        }
165    }
166
167    /// Save multiple key-value pairs concurrently.
168    ///
169    /// All writes are issued simultaneously using `try_join_all` instead of
170    /// the default sequential loop, so throughput scales with disk/OS concurrency.
171    async fn batch_save(&self, items: &[(&str, &[u8])]) -> Result<(), AgentRuntimeError> {
172        let futs: Vec<_> = items.iter().map(|(key, value)| self.save(key, value)).collect();
173        try_join_all(futs).await?;
174        Ok(())
175    }
176
177    /// Load multiple keys concurrently.
178    ///
179    /// All reads are issued simultaneously using `try_join_all`.
180    /// Returns results in the same order as `keys`.
181    async fn batch_load(&self, keys: &[&str]) -> Result<Vec<Option<Vec<u8>>>, AgentRuntimeError> {
182        let futs: Vec<_> = keys.iter().map(|key| self.load(key)).collect();
183        try_join_all(futs).await
184    }
185}
186
187impl FilePersistenceBackend {
188    /// List all keys currently stored in the backend.
189    ///
190    /// Returns the readable prefix portion of stored filenames (the part before
191    /// the hash suffix). Useful for backup/restore tooling and debugging.
192    pub async fn list_keys(&self) -> Result<Vec<String>, AgentRuntimeError> {
193        let mut entries = tokio::fs::read_dir(self.base_dir.as_ref())
194            .await
195            .map_err(|e| AgentRuntimeError::Persistence(format!("list_keys readdir: {e}")))?;
196
197        let mut keys = Vec::new();
198        while let Some(entry) = entries
199            .next_entry()
200            .await
201            .map_err(|e| AgentRuntimeError::Persistence(format!("list_keys entry: {e}")))?
202        {
203            if let Some(name) = entry.file_name().to_str() {
204                if name.ends_with(".bin") {
205                    if let Some(stem) = name.strip_suffix(".bin") {
206                        // Strip the "-<16hexdigits>" suffix to recover the readable prefix.
207                        if stem.len() > 17 {
208                            let prefix = &stem[..stem.len() - 17]; // 16 hex + 1 dash
209                            keys.push(prefix.to_owned());
210                        }
211                    }
212                }
213            }
214        }
215        Ok(keys)
216    }
217
218    /// Check whether a key exists in the backend without loading its value.
219    ///
220    /// Returns `Ok(true)` if the file for `key` is present in the store.
221    pub async fn exists(&self, key: &str) -> Result<bool, AgentRuntimeError> {
222        let path = self.path_for(key);
223        Ok(tokio::fs::metadata(&path).await.is_ok())
224    }
225
226    /// Return the number of keys currently stored in the backend.
227    ///
228    /// More efficient than `list_keys().await?.len()` for callers that only
229    /// need the count, as it avoids building the string list.
230    pub async fn key_count(&self) -> Result<usize, AgentRuntimeError> {
231        let mut entries = tokio::fs::read_dir(self.base_dir.as_ref())
232            .await
233            .map_err(|e| AgentRuntimeError::Persistence(format!("key_count readdir: {e}")))?;
234        let mut count = 0usize;
235        while let Some(entry) = entries
236            .next_entry()
237            .await
238            .map_err(|e| AgentRuntimeError::Persistence(format!("key_count entry: {e}")))?
239        {
240            if entry
241                .file_name()
242                .to_str()
243                .map_or(false, |n| n.ends_with(".bin"))
244            {
245                count += 1;
246            }
247        }
248        Ok(count)
249    }
250}
251
252// ── Tests ─────────────────────────────────────────────────────────────────────
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use std::sync::Arc;
258
259    async fn temp_backend() -> (FilePersistenceBackend, tempdir::Handle) {
260        // Fall back to a simple temp path approach without the tempdir crate.
261        // We use tokio's temp file support via std's temp_dir.
262        let dir = std::env::temp_dir().join(format!("agent_runtime_test_{}", uuid::Uuid::new_v4()));
263        tokio::fs::create_dir_all(&dir).await.unwrap();
264        let backend = FilePersistenceBackend::new(dir.clone());
265        (backend, tempdir::Handle { path: dir })
266    }
267
268    // Use a simple inline struct instead of a crate dependency.
269    mod tempdir {
270        pub struct Handle {
271            pub path: std::path::PathBuf,
272        }
273        impl Drop for Handle {
274            fn drop(&mut self) {
275                let _ = std::fs::remove_dir_all(&self.path);
276            }
277        }
278    }
279
280    #[tokio::test]
281    async fn test_file_backend_save_and_load() {
282        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
283        tokio::fs::create_dir_all(&dir).await.unwrap();
284        let _guard = tempdir::Handle { path: dir.clone() };
285        let backend = FilePersistenceBackend::new(&dir);
286
287        backend.save("test-key", b"hello world").await.unwrap();
288        let loaded = backend.load("test-key").await.unwrap();
289        assert_eq!(loaded, Some(b"hello world".to_vec()));
290    }
291
292    #[tokio::test]
293    async fn test_file_backend_load_missing_returns_none() {
294        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
295        tokio::fs::create_dir_all(&dir).await.unwrap();
296        let _guard = tempdir::Handle { path: dir.clone() };
297        let backend = FilePersistenceBackend::new(&dir);
298
299        let loaded = backend.load("nonexistent").await.unwrap();
300        assert_eq!(loaded, None);
301    }
302
303    #[tokio::test]
304    async fn test_file_backend_delete_removes_file() {
305        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
306        tokio::fs::create_dir_all(&dir).await.unwrap();
307        let _guard = tempdir::Handle { path: dir.clone() };
308        let backend = FilePersistenceBackend::new(&dir);
309
310        backend.save("to-delete", b"data").await.unwrap();
311        backend.delete("to-delete").await.unwrap();
312        let loaded = backend.load("to-delete").await.unwrap();
313        assert_eq!(loaded, None);
314    }
315
316    #[tokio::test]
317    async fn test_file_backend_delete_missing_is_noop() {
318        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
319        tokio::fs::create_dir_all(&dir).await.unwrap();
320        let _guard = tempdir::Handle { path: dir.clone() };
321        let backend = FilePersistenceBackend::new(&dir);
322
323        // Should not error
324        backend.delete("never-existed").await.unwrap();
325    }
326
327    #[tokio::test]
328    async fn test_file_backend_key_sanitization() {
329        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
330        tokio::fs::create_dir_all(&dir).await.unwrap();
331        let _guard = tempdir::Handle { path: dir.clone() };
332        let backend = FilePersistenceBackend::new(&dir);
333
334        // Key with path separators must not escape base_dir
335        backend.save("agent/session:1", b"data").await.unwrap();
336        let loaded = backend.load("agent/session:1").await.unwrap();
337        assert_eq!(loaded, Some(b"data".to_vec()));
338    }
339
340    #[tokio::test]
341    async fn test_file_backend_collision_resistant_keys() {
342        // "a/b" and "a_b" both sanitize to "a_b" with naive replacement.
343        // The hash suffix must differentiate them.
344        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
345        tokio::fs::create_dir_all(&dir).await.unwrap();
346        let _guard = tempdir::Handle { path: dir.clone() };
347        let backend = FilePersistenceBackend::new(&dir);
348
349        backend.save("a/b", b"slash").await.unwrap();
350        backend.save("a_b", b"underscore").await.unwrap();
351
352        let loaded_slash = backend.load("a/b").await.unwrap();
353        let loaded_under = backend.load("a_b").await.unwrap();
354        assert_eq!(loaded_slash, Some(b"slash".to_vec()));
355        assert_eq!(loaded_under, Some(b"underscore".to_vec()));
356        assert_ne!(
357            loaded_slash, loaded_under,
358            "keys with the same sanitized form must not collide"
359        );
360    }
361
362    #[tokio::test]
363    async fn test_batch_save_and_load() {
364        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
365        tokio::fs::create_dir_all(&dir).await.unwrap();
366        let _guard = tempdir::Handle { path: dir.clone() };
367        let backend = FilePersistenceBackend::new(&dir);
368
369        let data: Vec<(&str, Vec<u8>)> = vec![
370            ("batch-key-1", b"value1".to_vec()),
371            ("batch-key-2", b"value2".to_vec()),
372        ];
373        let refs: Vec<(&str, &[u8])> = data.iter().map(|(k, v)| (*k, v.as_slice())).collect();
374        backend.batch_save(&refs).await.unwrap();
375
376        let keys = vec!["batch-key-1", "batch-key-2", "batch-key-missing"];
377        let results = backend.batch_load(&keys).await.unwrap();
378        assert_eq!(results[0], Some(b"value1".to_vec()));
379        assert_eq!(results[1], Some(b"value2".to_vec()));
380        assert_eq!(results[2], None);
381    }
382
383    #[tokio::test]
384    async fn test_file_backend_list_keys() {
385        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
386        tokio::fs::create_dir_all(&dir).await.unwrap();
387        let _guard = tempdir::Handle { path: dir.clone() };
388        let backend = FilePersistenceBackend::new(&dir);
389
390        backend.save("my-session", b"data1").await.unwrap();
391        backend.save("another-key", b"data2").await.unwrap();
392
393        let keys = backend.list_keys().await.unwrap();
394        assert_eq!(keys.len(), 2);
395    }
396
397    #[tokio::test]
398    async fn test_persistence_backend_object_safe() {
399        // Verify FilePersistenceBackend can be used as a trait object.
400        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
401        tokio::fs::create_dir_all(&dir).await.unwrap();
402        let _guard = tempdir::Handle { path: dir.clone() };
403        let backend: Arc<dyn PersistenceBackend> = Arc::new(FilePersistenceBackend::new(&dir));
404        backend.save("obj-safe", b"ok").await.unwrap();
405        let r = backend.load("obj-safe").await.unwrap();
406        assert_eq!(r, Some(b"ok".to_vec()));
407    }
408
409    // ── Round 11: key_count ───────────────────────────────────────────────────
410
411    #[tokio::test]
412    async fn test_key_count_zero_for_empty_directory() {
413        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
414        tokio::fs::create_dir_all(&dir).await.unwrap();
415        let _guard = tempdir::Handle { path: dir.clone() };
416        let backend = FilePersistenceBackend::new(&dir);
417        assert_eq!(backend.key_count().await.unwrap(), 0);
418    }
419
420    #[tokio::test]
421    async fn test_key_count_matches_number_of_saves() {
422        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
423        tokio::fs::create_dir_all(&dir).await.unwrap();
424        let _guard = tempdir::Handle { path: dir.clone() };
425        let backend = FilePersistenceBackend::new(&dir);
426        backend.save("k1", b"v1").await.unwrap();
427        backend.save("k2", b"v2").await.unwrap();
428        backend.save("k3", b"v3").await.unwrap();
429        assert_eq!(backend.key_count().await.unwrap(), 3);
430    }
431
432    #[tokio::test]
433    async fn test_key_count_decrements_after_delete() {
434        let dir = std::env::temp_dir().join(format!("art_{}", uuid::Uuid::new_v4()));
435        tokio::fs::create_dir_all(&dir).await.unwrap();
436        let _guard = tempdir::Handle { path: dir.clone() };
437        let backend = FilePersistenceBackend::new(&dir);
438        backend.save("key", b"val").await.unwrap();
439        assert_eq!(backend.key_count().await.unwrap(), 1);
440        backend.delete("key").await.unwrap();
441        assert_eq!(backend.key_count().await.unwrap(), 0);
442    }
443}
llm_agent_runtime/persistence.rs

llm_agent_runtime/
persistence.rs
