mockforge_vbr/
session.rs

1//! Session integration
2//!
3//! This module integrates with existing SessionManager from mockforge-core/intelligent_behavior
4//! for session-scoped data and session expiration handling.
5
6use crate::config::StorageBackend;
7use crate::database::{create_database, VirtualDatabase};
8use crate::Result;
9use mockforge_core::intelligent_behavior::session::SessionManager;
10use std::collections::HashMap;
11use std::sync::Arc;
12use tokio::sync::RwLock;
13
14/// Session-scoped data manager
15///
16/// Manages per-session virtual databases when session-scoped data is enabled.
17/// Each session gets its own isolated database instance.
18pub struct SessionDataManager {
19    /// Session manager from mockforge-core
20    pub session_manager: Arc<SessionManager>,
21    /// Storage backend configuration for session databases
22    storage_backend: StorageBackend,
23    /// Per-session database instances (in-memory by default)
24    session_databases: Arc<RwLock<HashMap<String, Arc<dyn VirtualDatabase + Send + Sync>>>>,
25}
26
27impl SessionDataManager {
28    /// Create a new session data manager
29    pub fn new(session_manager: Arc<SessionManager>, storage_backend: StorageBackend) -> Self {
30        Self {
31            session_manager,
32            storage_backend,
33            session_databases: Arc::new(RwLock::new(HashMap::new())),
34        }
35    }
36
37    /// Get or create session-scoped database instance
38    ///
39    /// If session-scoped data is enabled, each session gets its own isolated
40    /// database. This allows different users/sessions to have separate data.
41    pub async fn get_session_database(
42        &self,
43        session_id: &str,
44    ) -> Result<Arc<dyn VirtualDatabase + Send + Sync>> {
45        // Check if session exists in SessionManager
46        let session_state = self.session_manager.get_session(session_id).await;
47        if session_state.is_none() {
48            return Err(crate::Error::generic(format!("Session '{}' not found", session_id)));
49        }
50
51        // Check if we already have a database for this session
52        let databases = self.session_databases.read().await;
53        if let Some(db) = databases.get(session_id) {
54            return Ok(Arc::clone(db));
55        }
56        drop(databases);
57
58        // Create a new in-memory database for this session
59        // Sessions always use in-memory storage for isolation
60        let db_arc = create_database(&StorageBackend::Memory).await?;
61        // Note: initialize() is called during create_database, so we don't need to call it again
62        let db_clone = Arc::clone(&db_arc);
63
64        // Store it for future use
65        let mut databases = self.session_databases.write().await;
66        databases.insert(session_id.to_string(), db_arc);
67
68        Ok(db_clone)
69    }
70
71    /// Clean up database for a session
72    pub async fn cleanup_session_database(&self, session_id: &str) -> Result<()> {
73        let mut databases = self.session_databases.write().await;
74        databases.remove(session_id);
75        // Database will be dropped when Arc is dropped
76        Ok(())
77    }
78
79    /// Clean up databases for expired sessions
80    pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
81        // Get expired sessions from SessionManager
82        let expired_count = self.session_manager.cleanup_expired_sessions().await;
83
84        // Get list of active sessions
85        let active_sessions = self.session_manager.list_sessions().await;
86        let active_set: std::collections::HashSet<String> = active_sessions.into_iter().collect();
87
88        // Remove databases for sessions that are no longer active
89        let mut databases = self.session_databases.write().await;
90        let mut removed = 0;
91        let expired_db_sessions: Vec<String> = databases
92            .keys()
93            .filter(|session_id| !active_set.contains(*session_id))
94            .cloned()
95            .collect();
96
97        for session_id in expired_db_sessions {
98            databases.remove(&session_id);
99            removed += 1;
100        }
101
102        Ok(removed)
103    }
104
105    /// Get the session manager
106    pub fn session_manager(&self) -> &Arc<SessionManager> {
107        &self.session_manager
108    }
109
110    /// Get session state (returns session if it exists)
111    pub async fn get_session_state(
112        &self,
113        session_id: &str,
114    ) -> Option<mockforge_core::intelligent_behavior::types::SessionState> {
115        self.session_manager.get_session(session_id).await
116    }
117
118    /// Update session state
119    pub async fn update_session_state(
120        &self,
121        session_id: &str,
122        state: mockforge_core::intelligent_behavior::types::SessionState,
123    ) -> Result<()> {
124        self.session_manager.update_session(session_id, state).await
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use mockforge_core::intelligent_behavior::session::SessionTracking;
132
133    fn create_test_session_manager() -> Arc<SessionManager> {
134        let config = SessionTracking::default();
135        Arc::new(SessionManager::new(config, 3600))
136    }
137
138    #[tokio::test]
139    async fn test_session_data_manager_new() {
140        let session_manager = create_test_session_manager();
141        let manager = SessionDataManager::new(session_manager.clone(), StorageBackend::Memory);
142
143        assert!(Arc::ptr_eq(&manager.session_manager, &session_manager));
144    }
145
146    #[tokio::test]
147    async fn test_get_session_database_no_session() {
148        let session_manager = create_test_session_manager();
149        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
150
151        let result = manager.get_session_database("nonexistent-session").await;
152        assert!(result.is_err());
153        match result {
154            Err(err) => {
155                assert!(err.to_string().contains("Session"));
156                assert!(err.to_string().contains("not found"));
157            }
158            Ok(_) => panic!("Expected an error"),
159        }
160    }
161
162    #[tokio::test]
163    async fn test_get_session_database_creates_new() {
164        let session_manager = create_test_session_manager();
165
166        // Create a session first
167        let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
168
169        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
170        let result = manager.get_session_database("test-session").await;
171
172        assert!(result.is_ok());
173    }
174
175    #[tokio::test]
176    async fn test_get_session_database_returns_same() {
177        let session_manager = create_test_session_manager();
178
179        // Create a session first
180        let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
181
182        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
183
184        // Get database twice
185        let db1 = manager.get_session_database("test-session").await.unwrap();
186        let db2 = manager.get_session_database("test-session").await.unwrap();
187
188        // Should return the same Arc
189        assert!(Arc::ptr_eq(&db1, &db2));
190    }
191
192    #[tokio::test]
193    async fn test_cleanup_session_database() {
194        let session_manager = create_test_session_manager();
195
196        // Create a session first
197        let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
198
199        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
200
201        // Get a database to create the entry
202        let _db = manager.get_session_database("test-session").await.unwrap();
203
204        // Clean up
205        let result = manager.cleanup_session_database("test-session").await;
206        assert!(result.is_ok());
207    }
208
209    #[tokio::test]
210    async fn test_cleanup_nonexistent_session() {
211        let session_manager = create_test_session_manager();
212        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
213
214        // Should not error even if session doesn't exist
215        let result = manager.cleanup_session_database("nonexistent").await;
216        assert!(result.is_ok());
217    }
218
219    #[tokio::test]
220    async fn test_session_manager_getter() {
221        let session_manager = create_test_session_manager();
222        let manager = SessionDataManager::new(session_manager.clone(), StorageBackend::Memory);
223
224        assert!(Arc::ptr_eq(manager.session_manager(), &session_manager));
225    }
226
227    #[tokio::test]
228    async fn test_get_session_state_exists() {
229        let session_manager = create_test_session_manager();
230
231        // Create a session first
232        let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
233
234        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
235        let state = manager.get_session_state("test-session").await;
236
237        assert!(state.is_some());
238    }
239
240    #[tokio::test]
241    async fn test_get_session_state_not_exists() {
242        let session_manager = create_test_session_manager();
243        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
244
245        let state = manager.get_session_state("nonexistent").await;
246        assert!(state.is_none());
247    }
248
249    #[tokio::test]
250    async fn test_multiple_sessions_isolated() {
251        let session_manager = create_test_session_manager();
252
253        // Create two sessions
254        let _ = session_manager.get_or_create_session(Some("session-1".to_string())).await;
255        let _ = session_manager.get_or_create_session(Some("session-2".to_string())).await;
256
257        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
258
259        let db1 = manager.get_session_database("session-1").await.unwrap();
260        let db2 = manager.get_session_database("session-2").await.unwrap();
261
262        // Should be different databases
263        assert!(!Arc::ptr_eq(&db1, &db2));
264    }
265
266    #[tokio::test]
267    async fn test_cleanup_expired_sessions() {
268        let session_manager = create_test_session_manager();
269        let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
270
271        let result = manager.cleanup_expired_sessions().await;
272        assert!(result.is_ok());
273    }
274}
mockforge_vbr/session.rs

mockforge_vbr/
session.rs
