ai_agents_core/traits/
storage.rs

1//! Storage trait for agent persistence
2
3use async_trait::async_trait;
4
5use crate::error::{AgentError, Result};
6use crate::types::{
7    FactFilter, KeyFact, SessionFilter, SessionMetadata, SessionSummary, StateMachineSnapshot,
8};
9
10/// Minimal record of a spawned agent for session persistence.
11#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
12pub struct SpawnedAgentEntry {
13    pub id: String,
14    pub name: String,
15    pub spec_yaml: String,
16}
17
18/// Snapshot of agent state for persistence
19#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
20pub struct AgentSnapshot {
21    pub version: String,
22    pub agent_id: String,
23    pub timestamp: chrono::DateTime<chrono::Utc>,
24    #[serde(default)]
25    pub state_machine: Option<StateMachineSnapshot>,
26    pub memory: super::memory::MemorySnapshot,
27    #[serde(default)]
28    pub context: std::collections::HashMap<String, serde_json::Value>,
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub spawned_agents: Option<Vec<SpawnedAgentEntry>>,
31    /// Persona snapshot (serialized as Value to avoid core->persona dependency).
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub persona: Option<serde_json::Value>,
34    /// Relationship snapshot (serialized as Value to avoid core->relationships dependency).
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub relationships: Option<serde_json::Value>,
37}
38
39impl AgentSnapshot {
40    pub fn new(agent_id: String) -> Self {
41        Self {
42            version: env!("CARGO_PKG_VERSION").to_string(),
43            agent_id,
44            timestamp: chrono::Utc::now(),
45            state_machine: None,
46            memory: super::memory::MemorySnapshot::default(),
47            context: std::collections::HashMap::new(),
48            spawned_agents: None,
49            persona: None,
50            relationships: None,
51        }
52    }
53
54    pub fn with_state_machine(mut self, snapshot: StateMachineSnapshot) -> Self {
55        self.state_machine = Some(snapshot);
56        self
57    }
58
59    pub fn with_memory(mut self, snapshot: super::memory::MemorySnapshot) -> Self {
60        self.memory = snapshot;
61        self
62    }
63
64    pub fn with_context(
65        mut self,
66        context: std::collections::HashMap<String, serde_json::Value>,
67    ) -> Self {
68        self.context = context;
69        self
70    }
71
72    pub fn with_spawned_agents(mut self, agents: Vec<SpawnedAgentEntry>) -> Self {
73        self.spawned_agents = Some(agents);
74        self
75    }
76}
77
78/// An optional storage feature that a backend supports.
79#[non_exhaustive]
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
81pub enum StorageCapability {
82    Snapshot,
83    SessionMetadata,
84    SessionFiltering,
85    ExpiryCleanup,
86    ActorFacts,
87    ActorRelationships,
88    ActorDataDeletion,
89}
90
91impl std::fmt::Display for StorageCapability {
92    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93        match self {
94            Self::Snapshot => f.write_str("snapshot"),
95            Self::SessionMetadata => f.write_str("session metadata"),
96            Self::SessionFiltering => f.write_str("session filtering"),
97            Self::ExpiryCleanup => f.write_str("expiry cleanup"),
98            Self::ActorFacts => f.write_str("actor facts"),
99            Self::ActorRelationships => f.write_str("actor relationships"),
100            Self::ActorDataDeletion => f.write_str("atomic actor data deletion"),
101        }
102    }
103}
104
105/// Core storage trait for persisting agent state.
106///
107/// Built-in backends: `FileStorage`, `SqliteStorage`, and `RedisStorage`.
108/// Implement this for custom persistence (e.g., PostgreSQL, DynamoDB).
109#[async_trait]
110pub trait AgentStorage: Send + Sync {
111    /// Return whether this backend implements a storage feature.
112    fn supports(&self, capability: StorageCapability) -> bool {
113        capability == StorageCapability::Snapshot
114    }
115
116    /// Persist an agent snapshot for the given session ID.
117    async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()>;
118    /// Load a snapshot. Returns `None` if the session does not exist.
119    async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>>;
120    /// Remove a session's persisted data.
121    async fn delete(&self, session_id: &str) -> Result<()>;
122    /// List all stored session IDs.
123    async fn list_sessions(&self) -> Result<Vec<String>>;
124
125    // --- Session metadata ---
126
127    /// Atomically persist a snapshot and its metadata when the backend supports session metadata.
128    async fn save_snapshot_with_metadata(
129        &self,
130        session_id: &str,
131        snapshot: &AgentSnapshot,
132        metadata: &SessionMetadata,
133    ) -> Result<()> {
134        let _ = (session_id, snapshot, metadata);
135        Err(AgentError::UnsupportedStorageCapability(
136            StorageCapability::SessionMetadata,
137        ))
138    }
139
140    /// Save session metadata alongside the snapshot.
141    async fn save_metadata(&self, _session_id: &str, _metadata: &SessionMetadata) -> Result<()> {
142        Err(AgentError::UnsupportedStorageCapability(
143            StorageCapability::SessionMetadata,
144        ))
145    }
146
147    /// Load session metadata.
148    async fn load_metadata(&self, _session_id: &str) -> Result<Option<SessionMetadata>> {
149        Err(AgentError::UnsupportedStorageCapability(
150            StorageCapability::SessionMetadata,
151        ))
152    }
153
154    /// List sessions matching a filter.
155    async fn list_sessions_filtered(&self, _filter: &SessionFilter) -> Result<Vec<SessionSummary>> {
156        Err(AgentError::UnsupportedStorageCapability(
157            StorageCapability::SessionFiltering,
158        ))
159    }
160
161    /// Delete sessions that have expired based on TTL.
162    async fn cleanup_expired(&self) -> Result<usize> {
163        Err(AgentError::UnsupportedStorageCapability(
164            StorageCapability::ExpiryCleanup,
165        ))
166    }
167
168    // --- Actor facts ---
169
170    /// Persist facts for an actor. Merges with existing facts.
171    async fn save_facts(&self, _agent_id: &str, _actor_id: &str, _facts: &[KeyFact]) -> Result<()> {
172        Err(AgentError::UnsupportedStorageCapability(
173            StorageCapability::ActorFacts,
174        ))
175    }
176
177    /// Load all facts for a specific actor.
178    async fn load_facts(&self, _agent_id: &str, _actor_id: &str) -> Result<Vec<KeyFact>> {
179        Err(AgentError::UnsupportedStorageCapability(
180            StorageCapability::ActorFacts,
181        ))
182    }
183
184    /// Load facts matching a filter (cross-actor queries).
185    async fn query_facts(&self, _agent_id: &str, _filter: &FactFilter) -> Result<Vec<KeyFact>> {
186        Err(AgentError::UnsupportedStorageCapability(
187            StorageCapability::ActorFacts,
188        ))
189    }
190
191    /// Delete a single fact by ID.
192    async fn delete_fact(&self, _agent_id: &str, _actor_id: &str, _fact_id: &str) -> Result<()> {
193        Err(AgentError::UnsupportedStorageCapability(
194            StorageCapability::ActorFacts,
195        ))
196    }
197
198    /// Atomically delete facts, relationships, and sessions owned by an actor.
199    async fn delete_actor_data(&self, _agent_id: &str, _actor_id: &str) -> Result<()> {
200        Err(AgentError::UnsupportedStorageCapability(
201            StorageCapability::ActorDataDeletion,
202        ))
203    }
204
205    /// Persist relationship data for an actor. Value is owned by the relationships crate.
206    async fn save_relationship(
207        &self,
208        _agent_id: &str,
209        _actor_id: &str,
210        _relationship: &serde_json::Value,
211    ) -> Result<()> {
212        Err(AgentError::UnsupportedStorageCapability(
213            StorageCapability::ActorRelationships,
214        ))
215    }
216
217    /// Load relationship data for a specific actor.
218    async fn load_relationship(
219        &self,
220        _agent_id: &str,
221        _actor_id: &str,
222    ) -> Result<Option<serde_json::Value>> {
223        Err(AgentError::UnsupportedStorageCapability(
224            StorageCapability::ActorRelationships,
225        ))
226    }
227
228    /// List actor IDs that have relationship data for an agent.
229    async fn list_relationship_actors(&self, _agent_id: &str) -> Result<Vec<String>> {
230        Err(AgentError::UnsupportedStorageCapability(
231            StorageCapability::ActorRelationships,
232        ))
233    }
234
235    /// Delete relationship data for a specific actor.
236    async fn delete_relationship(&self, _agent_id: &str, _actor_id: &str) -> Result<()> {
237        Err(AgentError::UnsupportedStorageCapability(
238            StorageCapability::ActorRelationships,
239        ))
240    }
241}
242
243/// Storage backend for testing unsupported persistence behavior.
244pub struct NoopStorage;
245
246#[async_trait]
247impl AgentStorage for NoopStorage {
248    fn supports(&self, _capability: StorageCapability) -> bool {
249        false
250    }
251
252    async fn save(&self, _session_id: &str, _snapshot: &AgentSnapshot) -> Result<()> {
253        Err(AgentError::UnsupportedStorageCapability(
254            StorageCapability::Snapshot,
255        ))
256    }
257    async fn load(&self, _session_id: &str) -> Result<Option<AgentSnapshot>> {
258        Err(AgentError::UnsupportedStorageCapability(
259            StorageCapability::Snapshot,
260        ))
261    }
262    async fn delete(&self, _session_id: &str) -> Result<()> {
263        Err(AgentError::UnsupportedStorageCapability(
264            StorageCapability::Snapshot,
265        ))
266    }
267    async fn list_sessions(&self) -> Result<Vec<String>> {
268        Err(AgentError::UnsupportedStorageCapability(
269            StorageCapability::Snapshot,
270        ))
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277
278    struct SnapshotOnlyStorage;
279
280    struct TargetedActorStorage {
281        fact_deletes: std::sync::atomic::AtomicUsize,
282        relationship_deletes: std::sync::atomic::AtomicUsize,
283    }
284
285    #[async_trait]
286    impl AgentStorage for SnapshotOnlyStorage {
287        async fn save(&self, _session_id: &str, _snapshot: &AgentSnapshot) -> Result<()> {
288            Ok(())
289        }
290
291        async fn load(&self, _session_id: &str) -> Result<Option<AgentSnapshot>> {
292            Ok(None)
293        }
294
295        async fn delete(&self, _session_id: &str) -> Result<()> {
296            Ok(())
297        }
298
299        async fn list_sessions(&self) -> Result<Vec<String>> {
300            Ok(vec![])
301        }
302    }
303
304    #[async_trait]
305    impl AgentStorage for TargetedActorStorage {
306        fn supports(&self, capability: StorageCapability) -> bool {
307            matches!(
308                capability,
309                StorageCapability::Snapshot
310                    | StorageCapability::ActorFacts
311                    | StorageCapability::ActorRelationships
312            )
313        }
314
315        async fn save(&self, _session_id: &str, _snapshot: &AgentSnapshot) -> Result<()> {
316            Ok(())
317        }
318
319        async fn load(&self, _session_id: &str) -> Result<Option<AgentSnapshot>> {
320            Ok(None)
321        }
322
323        async fn delete(&self, _session_id: &str) -> Result<()> {
324            Ok(())
325        }
326
327        async fn list_sessions(&self) -> Result<Vec<String>> {
328            Ok(Vec::new())
329        }
330
331        async fn delete_fact(
332            &self,
333            _agent_id: &str,
334            _actor_id: &str,
335            _fact_id: &str,
336        ) -> Result<()> {
337            self.fact_deletes
338                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
339            Ok(())
340        }
341
342        async fn delete_relationship(&self, _agent_id: &str, _actor_id: &str) -> Result<()> {
343            self.relationship_deletes
344                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
345            Ok(())
346        }
347    }
348
349    fn assert_unsupported<T>(result: Result<T>, expected: StorageCapability) {
350        assert!(matches!(
351            result,
352            Err(AgentError::UnsupportedStorageCapability(capability)) if capability == expected
353        ));
354    }
355
356    #[test]
357    fn snapshot_only_custom_backend_uses_default_capability() {
358        let storage = SnapshotOnlyStorage;
359
360        assert!(storage.supports(StorageCapability::Snapshot));
361        assert!(!storage.supports(StorageCapability::SessionMetadata));
362        futures::executor::block_on(storage.save("session", &AgentSnapshot::new("agent".into())))
363            .unwrap();
364    }
365
366    #[test]
367    fn custom_backend_extensions_return_typed_errors() {
368        let storage = SnapshotOnlyStorage;
369        let metadata = SessionMetadata::default();
370        let session_filter = SessionFilter::default();
371        let fact_filter = FactFilter::default();
372        let relationship = serde_json::json!({});
373
374        assert_unsupported(
375            futures::executor::block_on(storage.save_snapshot_with_metadata(
376                "session",
377                &AgentSnapshot::new("agent".into()),
378                &metadata,
379            )),
380            StorageCapability::SessionMetadata,
381        );
382        assert_unsupported(
383            futures::executor::block_on(storage.save_metadata("session", &metadata)),
384            StorageCapability::SessionMetadata,
385        );
386        assert_unsupported(
387            futures::executor::block_on(storage.load_metadata("session")),
388            StorageCapability::SessionMetadata,
389        );
390        assert_unsupported(
391            futures::executor::block_on(storage.list_sessions_filtered(&session_filter)),
392            StorageCapability::SessionFiltering,
393        );
394        assert_unsupported(
395            futures::executor::block_on(storage.cleanup_expired()),
396            StorageCapability::ExpiryCleanup,
397        );
398        assert_unsupported(
399            futures::executor::block_on(storage.save_facts("agent", "actor", &[])),
400            StorageCapability::ActorFacts,
401        );
402        assert_unsupported(
403            futures::executor::block_on(storage.load_facts("agent", "actor")),
404            StorageCapability::ActorFacts,
405        );
406        assert_unsupported(
407            futures::executor::block_on(storage.query_facts("agent", &fact_filter)),
408            StorageCapability::ActorFacts,
409        );
410        assert_unsupported(
411            futures::executor::block_on(storage.delete_fact("agent", "actor", "fact")),
412            StorageCapability::ActorFacts,
413        );
414        assert_unsupported(
415            futures::executor::block_on(storage.delete_actor_data("agent", "actor")),
416            StorageCapability::ActorDataDeletion,
417        );
418        assert_unsupported(
419            futures::executor::block_on(storage.save_relationship("agent", "actor", &relationship)),
420            StorageCapability::ActorRelationships,
421        );
422        assert_unsupported(
423            futures::executor::block_on(storage.load_relationship("agent", "actor")),
424            StorageCapability::ActorRelationships,
425        );
426        assert_unsupported(
427            futures::executor::block_on(storage.list_relationship_actors("agent")),
428            StorageCapability::ActorRelationships,
429        );
430        assert_unsupported(
431            futures::executor::block_on(storage.delete_relationship("agent", "actor")),
432            StorageCapability::ActorRelationships,
433        );
434    }
435
436    #[test]
437    fn unsupported_composite_actor_deletion_does_not_call_targeted_deletes() {
438        let storage = TargetedActorStorage {
439            fact_deletes: std::sync::atomic::AtomicUsize::new(0),
440            relationship_deletes: std::sync::atomic::AtomicUsize::new(0),
441        };
442
443        assert_unsupported(
444            futures::executor::block_on(storage.delete_actor_data("agent", "actor")),
445            StorageCapability::ActorDataDeletion,
446        );
447        assert_eq!(
448            storage
449                .fact_deletes
450                .load(std::sync::atomic::Ordering::SeqCst),
451            0
452        );
453        assert_eq!(
454            storage
455                .relationship_deletes
456                .load(std::sync::atomic::Ordering::SeqCst),
457            0
458        );
459    }
460
461    #[test]
462    fn noop_storage_rejects_snapshot_operations() {
463        let storage = NoopStorage;
464        let snapshot = AgentSnapshot::new("agent".into());
465
466        assert!(!storage.supports(StorageCapability::Snapshot));
467        assert_unsupported(
468            futures::executor::block_on(storage.save("session", &snapshot)),
469            StorageCapability::Snapshot,
470        );
471        assert_unsupported(
472            futures::executor::block_on(storage.load("session")),
473            StorageCapability::Snapshot,
474        );
475        assert_unsupported(
476            futures::executor::block_on(storage.delete("session")),
477            StorageCapability::Snapshot,
478        );
479        assert_unsupported(
480            futures::executor::block_on(storage.list_sessions()),
481            StorageCapability::Snapshot,
482        );
483    }
484}
ai_agents_core/traits/storage.rs

ai_agents_core/traits/
storage.rs
