vex_persist/
context_store.rs

1//! Context packet storage
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use std::sync::Arc;
6use uuid::Uuid;
7
8use crate::backend::{StorageBackend, StorageError, StorageExt};
9use vex_core::ContextPacket;
10
11/// Serializable context state
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ContextState {
14    /// Unique ID for this context
15    pub id: Uuid,
16    /// The context packet
17    pub packet: ContextPacket,
18    /// Agent that created this context
19    pub agent_id: Option<Uuid>,
20    /// When it was stored
21    pub stored_at: DateTime<Utc>,
22}
23
24/// Context store for persistence
25#[derive(Debug)]
26pub struct ContextStore<B: StorageBackend + ?Sized> {
27    backend: Arc<B>,
28    prefix: String,
29}
30
31impl<B: StorageBackend + ?Sized> ContextStore<B> {
32    /// Create a new context store
33    pub fn new(backend: Arc<B>) -> Self {
34        Self {
35            backend,
36            prefix: "context:".to_string(),
37        }
38    }
39
40    fn key(&self, tenant_id: &str, id: Uuid) -> String {
41        format!("{}tenant:{}:{}", self.prefix, tenant_id, id)
42    }
43
44    fn agent_key(&self, tenant_id: &str, agent_id: Uuid) -> String {
45        format!("{}tenant:{}:agent:{}", self.prefix, tenant_id, agent_id)
46    }
47
48    /// Save a context packet
49    pub async fn save(
50        &self,
51        tenant_id: &str,
52        packet: &ContextPacket,
53    ) -> Result<Uuid, StorageError> {
54        let id = Uuid::new_v4();
55        let state = ContextState {
56            id,
57            packet: packet.clone(),
58            agent_id: packet.source_agent,
59            stored_at: Utc::now(),
60        };
61        self.backend.set(&self.key(tenant_id, id), &state).await?;
62
63        // Also index by agent if available
64        if let Some(agent_id) = packet.source_agent {
65            let mut agent_contexts: Vec<Uuid> = self
66                .backend
67                .get(&self.agent_key(tenant_id, agent_id))
68                .await?
69                .unwrap_or_default();
70            agent_contexts.push(id);
71            self.backend
72                .set(&self.agent_key(tenant_id, agent_id), &agent_contexts)
73                .await?;
74        }
75
76        Ok(id)
77    }
78
79    /// Load a context by ID
80    pub async fn load(
81        &self,
82        tenant_id: &str,
83        id: Uuid,
84    ) -> Result<Option<ContextPacket>, StorageError> {
85        let state: Option<ContextState> = self.backend.get(&self.key(tenant_id, id)).await?;
86        Ok(state.map(|s| s.packet))
87    }
88
89    /// Load all contexts for an agent
90    pub async fn load_by_agent(
91        &self,
92        tenant_id: &str,
93        agent_id: Uuid,
94    ) -> Result<Vec<ContextPacket>, StorageError> {
95        let context_ids: Vec<Uuid> = self
96            .backend
97            .get(&self.agent_key(tenant_id, agent_id))
98            .await?
99            .unwrap_or_default();
100
101        let mut contexts = Vec::new();
102        for id in context_ids {
103            if let Some(ctx) = self.load(tenant_id, id).await? {
104                contexts.push(ctx);
105            }
106        }
107        Ok(contexts)
108    }
109
110    /// Delete a context
111    pub async fn delete(&self, tenant_id: &str, id: Uuid) -> Result<bool, StorageError> {
112        self.backend.delete(&self.key(tenant_id, id)).await
113    }
114
115    /// Get total count of stored contexts for a tenant
116    pub async fn count(&self, tenant_id: &str) -> Result<usize, StorageError> {
117        let tenant_prefix = format!("{}tenant:{}:", self.prefix, tenant_id);
118        let keys = self.backend.list_keys(&tenant_prefix).await?;
119        Ok(keys.iter().filter(|k| !k.contains(":agent:")).count())
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::backend::MemoryBackend;
127
128    #[tokio::test]
129    async fn test_context_store() {
130        let backend = Arc::new(MemoryBackend::new());
131        let store = ContextStore::new(backend);
132        let tenant_id = "test-tenant";
133
134        let mut packet = ContextPacket::new("Test content");
135        let agent_id = Uuid::new_v4();
136        packet.source_agent = Some(agent_id);
137
138        // Save
139        let id = store.save(tenant_id, &packet).await.unwrap();
140
141        // Load
142        let loaded = store.load(tenant_id, id).await.unwrap().unwrap();
143        assert_eq!(loaded.content, "Test content");
144
145        // Load by agent
146        let agent_contexts = store.load_by_agent(tenant_id, agent_id).await.unwrap();
147        assert_eq!(agent_contexts.len(), 1);
148
149        // Count
150        assert_eq!(store.count(tenant_id).await.unwrap(), 1);
151    }
152}
vex_persist/context_store.rs

vex_persist/
context_store.rs
