vex_persist/
agent_store.rs

1//! Agent storage
2
3use serde::{Deserialize, Serialize};
4use std::sync::Arc;
5use uuid::Uuid;
6
7use crate::backend::{StorageBackend, StorageError, StorageExt};
8use vex_core::{Agent, AgentConfig};
9
10/// Serializable agent state
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct AgentState {
13    pub id: Uuid,
14    pub parent_id: Option<Uuid>,
15    pub config: AgentConfig,
16    pub generation: u32,
17    pub fitness: f64,
18    pub created_at: chrono::DateTime<chrono::Utc>,
19    pub updated_at: chrono::DateTime<chrono::Utc>,
20}
21
22impl From<&Agent> for AgentState {
23    fn from(agent: &Agent) -> Self {
24        Self {
25            id: agent.id,
26            parent_id: agent.parent_id,
27            config: agent.config.clone(),
28            generation: agent.generation,
29            fitness: agent.fitness,
30            created_at: chrono::Utc::now(),
31            updated_at: chrono::Utc::now(),
32        }
33    }
34}
35
36impl AgentState {
37    /// Convert back to an Agent
38    pub fn to_agent(&self) -> Agent {
39        let mut agent = Agent::new(self.config.clone());
40        agent.id = self.id;
41        agent.parent_id = self.parent_id;
42        agent.generation = self.generation;
43        agent.fitness = self.fitness;
44        agent
45    }
46}
47
48/// Agent store for persistence
49#[derive(Debug)]
50pub struct AgentStore<B: StorageBackend + ?Sized> {
51    backend: Arc<B>,
52    prefix: String,
53}
54
55impl<B: StorageBackend + ?Sized> AgentStore<B> {
56    /// Create a new agent store
57    pub fn new(backend: Arc<B>) -> Self {
58        Self {
59            backend,
60            prefix: "agent:".to_string(),
61        }
62    }
63
64    /// Create with custom prefix
65    pub fn with_prefix(backend: Arc<B>, prefix: &str) -> Self {
66        Self {
67            backend,
68            prefix: prefix.to_string(),
69        }
70    }
71
72    fn key(&self, tenant_id: &str, id: Uuid) -> String {
73        format!("{}tenant:{}:{}", self.prefix, tenant_id, id)
74    }
75
76    /// Save an agent
77    pub async fn save(&self, tenant_id: &str, agent: &Agent) -> Result<(), StorageError> {
78        let state = AgentState::from(agent);
79        self.backend
80            .set(&self.key(tenant_id, agent.id), &state)
81            .await
82    }
83
84    /// Load an agent by ID
85    pub async fn load(&self, tenant_id: &str, id: Uuid) -> Result<Option<Agent>, StorageError> {
86        let state: Option<AgentState> = self.backend.get(&self.key(tenant_id, id)).await?;
87        Ok(state.map(|s| s.to_agent()))
88    }
89
90    /// Delete an agent
91    pub async fn delete(&self, tenant_id: &str, id: Uuid) -> Result<bool, StorageError> {
92        self.backend.delete(&self.key(tenant_id, id)).await
93    }
94
95    /// Check if agent exists
96    pub async fn exists(&self, tenant_id: &str, id: Uuid) -> Result<bool, StorageError> {
97        self.backend.exists(&self.key(tenant_id, id)).await
98    }
99
100    /// List all agent IDs for a tenant
101    pub async fn list(&self, tenant_id: &str) -> Result<Vec<Uuid>, StorageError> {
102        let tenant_prefix = format!("{}tenant:{}:", self.prefix, tenant_id);
103        let keys = self.backend.list_keys(&tenant_prefix).await?;
104        let ids: Vec<Uuid> = keys
105            .iter()
106            .filter_map(|k| {
107                k.strip_prefix(&tenant_prefix)
108                    .and_then(|s| Uuid::parse_str(s).ok())
109            })
110            .collect();
111        Ok(ids)
112    }
113
114    /// Load all agents for a tenant
115    pub async fn load_all(&self, tenant_id: &str) -> Result<Vec<Agent>, StorageError> {
116        let ids = self.list(tenant_id).await?;
117        let mut agents = Vec::new();
118        for id in ids {
119            if let Some(agent) = self.load(tenant_id, id).await? {
120                agents.push(agent);
121            }
122        }
123        Ok(agents)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use crate::backend::MemoryBackend;
131
132    #[tokio::test]
133    async fn test_agent_store() {
134        let backend = Arc::new(MemoryBackend::new());
135        let store = AgentStore::new(backend);
136        let tenant_id = "test-tenant";
137
138        let agent = Agent::new(AgentConfig {
139            name: "TestAgent".to_string(),
140            role: "Tester".to_string(),
141            max_depth: 2,
142            spawn_shadow: true,
143        });
144        let id = agent.id;
145
146        // Save
147        store.save(tenant_id, &agent).await.unwrap();
148
149        // Exists
150        assert!(store.exists(tenant_id, id).await.unwrap());
151
152        // Load
153        let loaded = store.load(tenant_id, id).await.unwrap().unwrap();
154        assert_eq!(loaded.id, id);
155        assert_eq!(loaded.config.name, "TestAgent");
156
157        // List
158        let ids = store.list(tenant_id).await.unwrap();
159        assert_eq!(ids.len(), 1);
160
161        // Delete
162        assert!(store.delete(tenant_id, id).await.unwrap());
163        assert!(!store.exists(tenant_id, id).await.unwrap());
164    }
165}
vex_persist/agent_store.rs

vex_persist/
agent_store.rs
