ultrafast_mcp_sequential_thinking/session/
mod.rs

1//! # Session Management Module
2//!
3//! Session management functionality for the sequential thinking system.
4//!
5//! This module provides session creation, management, and persistence
6//! capabilities for thinking sessions.
7
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::sync::Arc;
11use tokio::sync::RwLock;
12use uuid::Uuid;
13
14use crate::thinking::{ThinkingEngine, ThinkingProgress, ThinkingStats, ThoughtData};
15
16/// Session metadata
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct SessionMetadata {
19    /// Session title
20    pub title: String,
21    /// Session description
22    pub description: Option<String>,
23    /// Session tags
24    pub tags: Vec<String>,
25    /// Session priority
26    pub priority: SessionPriority,
27    /// Session status
28    pub status: SessionStatus,
29    /// Created timestamp
30    pub created_at: chrono::DateTime<chrono::Utc>,
31    /// Last modified timestamp
32    pub last_modified: chrono::DateTime<chrono::Utc>,
33    /// Expires at timestamp
34    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
35    /// Custom metadata
36    pub custom_data: HashMap<String, serde_json::Value>,
37}
38
39/// Session priority levels
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
41pub enum SessionPriority {
42    Low,
43    Normal,
44    High,
45    Critical,
46}
47
48/// Session status
49#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
50pub enum SessionStatus {
51    Active,
52    Paused,
53    Completed,
54    Cancelled,
55    Expired,
56}
57
58impl Default for SessionMetadata {
59    fn default() -> Self {
60        Self {
61            title: String::new(),
62            description: None,
63            tags: Vec::new(),
64            priority: SessionPriority::Normal,
65            status: SessionStatus::Active,
66            created_at: chrono::Utc::now(),
67            last_modified: chrono::Utc::now(),
68            expires_at: None,
69            custom_data: HashMap::new(),
70        }
71    }
72}
73
74/// A thinking session
75#[derive(Debug, Clone)]
76pub struct ThinkingSession {
77    /// Session ID
78    pub session_id: String,
79    /// Session metadata
80    pub metadata: SessionMetadata,
81    /// Thinking engine
82    pub engine: ThinkingEngine,
83    #[allow(dead_code)]
84    lock: Arc<RwLock<()>>,
85}
86
87impl ThinkingSession {
88    /// Create a new thinking session
89    pub fn new(session_id: String, title: String) -> Self {
90        let metadata = SessionMetadata {
91            title,
92            ..Default::default()
93        };
94
95        Self {
96            session_id,
97            metadata,
98            engine: ThinkingEngine::new(),
99            lock: Arc::new(RwLock::new(())),
100        }
101    }
102
103    /// Create a new thinking session with metadata
104    pub fn with_metadata(session_id: String, metadata: SessionMetadata) -> Self {
105        Self {
106            session_id,
107            metadata,
108            engine: ThinkingEngine::new(),
109            lock: Arc::new(RwLock::new(())),
110        }
111    }
112
113    /// Get session ID
114    pub fn id(&self) -> &str {
115        &self.session_id
116    }
117
118    /// Get session title
119    pub fn title(&self) -> &str {
120        &self.metadata.title
121    }
122
123    /// Get session status
124    pub fn status(&self) -> &SessionStatus {
125        &self.metadata.status
126    }
127
128    /// Set session status
129    pub fn set_status(&mut self, status: SessionStatus) {
130        self.metadata.status = status;
131        self.metadata.last_modified = chrono::Utc::now();
132    }
133
134    /// Get session priority
135    pub fn priority(&self) -> &SessionPriority {
136        &self.metadata.priority
137    }
138
139    /// Set session priority
140    pub fn set_priority(&mut self, priority: SessionPriority) {
141        self.metadata.priority = priority;
142        self.metadata.last_modified = chrono::Utc::now();
143    }
144
145    /// Add a tag to the session
146    pub fn add_tag(&mut self, tag: String) {
147        if !self.metadata.tags.contains(&tag) {
148            self.metadata.tags.push(tag);
149            self.metadata.last_modified = chrono::Utc::now();
150        }
151    }
152
153    /// Remove a tag from the session
154    pub fn remove_tag(&mut self, tag: &str) {
155        self.metadata.tags.retain(|t| t != tag);
156        self.metadata.last_modified = chrono::Utc::now();
157    }
158
159    /// Set custom metadata
160    pub fn set_custom_data(&mut self, key: String, value: serde_json::Value) {
161        self.metadata.custom_data.insert(key, value);
162        self.metadata.last_modified = chrono::Utc::now();
163    }
164
165    /// Get custom metadata
166    pub fn get_custom_data(&self, key: &str) -> Option<&serde_json::Value> {
167        self.metadata.custom_data.get(key)
168    }
169
170    /// Check if session is expired
171    pub fn is_expired(&self) -> bool {
172        if let Some(expires_at) = self.metadata.expires_at {
173            chrono::Utc::now() > expires_at
174        } else {
175            false
176        }
177    }
178
179    /// Check if session is active
180    pub fn is_active(&self) -> bool {
181        self.metadata.status == SessionStatus::Active && !self.is_expired()
182    }
183
184    /// Get session progress
185    pub fn get_progress(&self) -> ThinkingProgress {
186        self.engine.get_progress().clone()
187    }
188
189    /// Get session statistics
190    pub fn get_stats(&self) -> ThinkingStats {
191        self.engine.get_stats().clone()
192    }
193
194    /// Get all thoughts in the session
195    pub fn get_thoughts(&self) -> Vec<ThoughtData> {
196        self.engine.get_thoughts().to_vec()
197    }
198
199    /// Get session age
200    pub fn age(&self) -> chrono::Duration {
201        chrono::Utc::now() - self.metadata.created_at
202    }
203
204    /// Get session duration
205    pub fn duration(&self) -> chrono::Duration {
206        self.metadata.last_modified - self.metadata.created_at
207    }
208}
209
210/// Session manager for handling multiple sessions
211pub struct SessionManager {
212    /// Active sessions
213    sessions: Arc<RwLock<HashMap<String, ThinkingSession>>>,
214    /// Session configuration
215    config: SessionManagerConfig,
216    /// Statistics
217    stats: Arc<RwLock<SessionManagerStats>>,
218}
219
220/// Session manager configuration
221#[derive(Debug, Clone)]
222pub struct SessionManagerConfig {
223    /// Maximum number of active sessions
224    pub max_sessions: usize,
225    /// Session timeout in seconds
226    pub session_timeout: u64,
227    /// Whether to auto-cleanup expired sessions
228    pub auto_cleanup: bool,
229    /// Cleanup interval in seconds
230    pub cleanup_interval: u64,
231    /// Whether to persist sessions
232    pub persist_sessions: bool,
233    /// Persistence directory
234    pub persistence_dir: String,
235}
236
237impl Default for SessionManagerConfig {
238    fn default() -> Self {
239        Self {
240            max_sessions: 100,
241            session_timeout: 3600,
242            auto_cleanup: true,
243            cleanup_interval: 300,
244            persist_sessions: false,
245            persistence_dir: "./sessions".to_string(),
246        }
247    }
248}
249
250/// Session manager statistics
251#[derive(Debug, Clone, Default)]
252pub struct SessionManagerStats {
253    /// Total sessions created
254    pub total_sessions_created: u64,
255    /// Total sessions completed
256    pub total_sessions_completed: u64,
257    /// Total sessions cancelled
258    pub total_sessions_cancelled: u64,
259    /// Total sessions expired
260    pub total_sessions_expired: u64,
261    /// Current active sessions
262    pub active_sessions: usize,
263    /// Average session duration in seconds
264    pub avg_session_duration: f64,
265    /// Total session time in seconds
266    pub total_session_time: u64,
267}
268
269impl SessionManager {
270    /// Create a new session manager
271    pub fn new() -> Self {
272        Self {
273            sessions: Arc::new(RwLock::new(HashMap::new())),
274            config: SessionManagerConfig::default(),
275            stats: Arc::new(RwLock::new(SessionManagerStats::default())),
276        }
277    }
278
279    /// Create a new session manager with configuration
280    pub fn with_config(config: SessionManagerConfig) -> Self {
281        Self {
282            sessions: Arc::new(RwLock::new(HashMap::new())),
283            config,
284            stats: Arc::new(RwLock::new(SessionManagerStats::default())),
285        }
286    }
287
288    /// Create a new session
289    pub async fn create_session(
290        &self,
291        title: String,
292    ) -> Result<String, Box<dyn std::error::Error>> {
293        let session_id = Uuid::new_v4().to_string();
294
295        // Check if we've reached the maximum number of sessions
296        {
297            let sessions = self.sessions.read().await;
298            if sessions.len() >= self.config.max_sessions {
299                return Err("Maximum number of sessions reached".into());
300            }
301        }
302
303        let session = ThinkingSession::new(session_id.clone(), title);
304
305        {
306            let mut sessions = self.sessions.write().await;
307            sessions.insert(session_id.clone(), session);
308        }
309
310        // Update statistics
311        {
312            let mut stats = self.stats.write().await;
313            stats.total_sessions_created += 1;
314            stats.active_sessions += 1;
315        }
316
317        Ok(session_id)
318    }
319
320    /// Get a session by ID
321    pub async fn get_session(&self, session_id: &str) -> Option<ThinkingSession> {
322        let sessions = self.sessions.read().await;
323        sessions.get(session_id).cloned()
324    }
325
326    /// Update a session
327    pub async fn update_session(&self, session_id: &str, session: ThinkingSession) -> bool {
328        let mut sessions = self.sessions.write().await;
329        sessions.insert(session_id.to_string(), session).is_some()
330    }
331
332    /// Remove a session
333    pub async fn remove_session(&self, session_id: &str) -> bool {
334        let mut sessions = self.sessions.write().await;
335        if sessions.remove(session_id).is_some() {
336            // Update statistics
337            let mut stats = self.stats.write().await;
338            stats.active_sessions = stats.active_sessions.saturating_sub(1);
339            true
340        } else {
341            false
342        }
343    }
344
345    /// List all session IDs
346    pub async fn list_session_ids(&self) -> Vec<String> {
347        let sessions = self.sessions.read().await;
348        sessions.keys().cloned().collect()
349    }
350
351    /// List active sessions
352    pub async fn list_active_sessions(&self) -> Vec<ThinkingSession> {
353        let sessions = self.sessions.read().await;
354        sessions
355            .values()
356            .filter(|session| session.is_active())
357            .cloned()
358            .collect()
359    }
360
361    /// Get session statistics
362    pub async fn get_stats(&self) -> SessionManagerStats {
363        self.stats.read().await.clone()
364    }
365
366    /// Cleanup expired sessions
367    pub async fn cleanup_expired_sessions(&self) -> usize {
368        let mut sessions = self.sessions.write().await;
369        let mut expired_count = 0;
370
371        let expired_sessions: Vec<String> = sessions
372            .iter()
373            .filter(|(_, session)| session.is_expired())
374            .map(|(id, _)| id.clone())
375            .collect();
376
377        for session_id in expired_sessions {
378            if let Some(session) = sessions.remove(&session_id) {
379                // Update statistics based on session status
380                let mut stats = self.stats.write().await;
381                match session.status() {
382                    SessionStatus::Completed => stats.total_sessions_completed += 1,
383                    SessionStatus::Cancelled => stats.total_sessions_cancelled += 1,
384                    _ => stats.total_sessions_expired += 1,
385                }
386                stats.active_sessions = stats.active_sessions.saturating_sub(1);
387                expired_count += 1;
388            }
389        }
390
391        expired_count
392    }
393
394    /// Start auto-cleanup task
395    pub async fn start_auto_cleanup(&self) {
396        let sessions = Arc::clone(&self.sessions);
397        let config = self.config.clone();
398        let stats = Arc::clone(&self.stats);
399
400        tokio::spawn(async move {
401            let mut interval =
402                tokio::time::interval(std::time::Duration::from_secs(config.cleanup_interval));
403
404            loop {
405                interval.tick().await;
406
407                let mut sessions_guard = sessions.write().await;
408                let mut expired_count = 0;
409
410                let expired_sessions: Vec<String> = sessions_guard
411                    .iter()
412                    .filter(|(_, session)| session.is_expired())
413                    .map(|(id, _)| id.clone())
414                    .collect();
415
416                for session_id in expired_sessions {
417                    if let Some(session) = sessions_guard.remove(&session_id) {
418                        // Update statistics
419                        let mut stats_guard = stats.write().await;
420                        match session.status() {
421                            SessionStatus::Completed => stats_guard.total_sessions_completed += 1,
422                            SessionStatus::Cancelled => stats_guard.total_sessions_cancelled += 1,
423                            _ => stats_guard.total_sessions_expired += 1,
424                        }
425                        stats_guard.active_sessions = stats_guard.active_sessions.saturating_sub(1);
426                        expired_count += 1;
427                    }
428                }
429
430                if expired_count > 0 {
431                    tracing::info!("Cleaned up {} expired sessions", expired_count);
432                }
433            }
434        });
435    }
436
437    /// Persist sessions to disk
438    pub async fn persist_sessions(&self) -> Result<(), Box<dyn std::error::Error>> {
439        if !self.config.persist_sessions {
440            return Ok(());
441        }
442
443        let sessions = self.sessions.read().await;
444        let sessions_data: HashMap<String, serde_json::Value> = sessions
445            .iter()
446            .map(|(id, session)| {
447                let session_data = serde_json::json!({
448                    "metadata": session.metadata,
449                    "thoughts": session.get_thoughts(),
450                    "stats": session.get_stats()
451                });
452                (id.clone(), session_data)
453            })
454            .collect();
455
456        let content = serde_json::to_string_pretty(&sessions_data)?;
457
458        // Ensure directory exists
459        std::fs::create_dir_all(&self.config.persistence_dir)?;
460
461        let file_path = format!("{}/sessions.json", self.config.persistence_dir);
462        std::fs::write(file_path, content)?;
463
464        Ok(())
465    }
466
467    /// Load sessions from disk
468    pub async fn load_sessions(&self) -> Result<(), Box<dyn std::error::Error>> {
469        if !self.config.persist_sessions {
470            return Ok(());
471        }
472
473        let file_path = format!("{}/sessions.json", self.config.persistence_dir);
474        if !std::path::Path::new(&file_path).exists() {
475            return Ok(());
476        }
477
478        let content = std::fs::read_to_string(file_path)?;
479        let sessions_data: HashMap<String, serde_json::Value> = serde_json::from_str(&content)?;
480
481        let mut sessions = self.sessions.write().await;
482        for (id, session_data) in sessions_data {
483            // Reconstruct session from persisted data
484            // This is a simplified implementation
485            let metadata: SessionMetadata = serde_json::from_value(
486                session_data
487                    .get("metadata")
488                    .unwrap_or(&serde_json::Value::Null)
489                    .clone(),
490            )?;
491
492            let session = ThinkingSession::with_metadata(id.clone(), metadata);
493            sessions.insert(id, session);
494        }
495
496        Ok(())
497    }
498}
499
500impl Default for SessionManager {
501    fn default() -> Self {
502        Self::new()
503    }
504}
505
506#[cfg(test)]
507mod tests {
508    use super::*;
509
510    #[test]
511    fn test_session_creation() {
512        let session = ThinkingSession::new("test-session".to_string(), "Test Session".to_string());
513        assert_eq!(session.id(), "test-session");
514        assert_eq!(session.title(), "Test Session");
515        assert_eq!(session.status(), &SessionStatus::Active);
516        assert_eq!(session.priority(), &SessionPriority::Normal);
517    }
518
519    #[test]
520    fn test_session_metadata() {
521        let mut session =
522            ThinkingSession::new("test-session".to_string(), "Test Session".to_string());
523
524        session.set_priority(SessionPriority::High);
525        assert_eq!(session.priority(), &SessionPriority::High);
526
527        session.add_tag("important".to_string());
528        assert!(session.metadata.tags.contains(&"important".to_string()));
529
530        session.set_custom_data("key".to_string(), serde_json::json!("value"));
531        assert_eq!(
532            session.get_custom_data("key"),
533            Some(&serde_json::json!("value"))
534        );
535    }
536
537    #[tokio::test]
538    async fn test_session_manager() {
539        let manager = SessionManager::new();
540
541        let session_id = manager
542            .create_session("Test Session".to_string())
543            .await
544            .unwrap();
545        assert!(!session_id.is_empty());
546
547        let session = manager.get_session(&session_id).await;
548        assert!(session.is_some());
549
550        let session_ids = manager.list_session_ids().await;
551        assert_eq!(session_ids.len(), 1);
552        assert!(session_ids.contains(&session_id));
553    }
554
555    #[tokio::test]
556    async fn test_session_cleanup() {
557        let manager = SessionManager::new();
558
559        // Create a session
560        let session_id = manager
561            .create_session("Test Session".to_string())
562            .await
563            .unwrap();
564
565        // Mark session as expired
566        if let Some(mut session) = manager.get_session(&session_id).await {
567            session.metadata.expires_at = Some(chrono::Utc::now() - chrono::Duration::hours(1));
568            manager.update_session(&session_id, session).await;
569        }
570
571        // Cleanup expired sessions
572        let expired_count = manager.cleanup_expired_sessions().await;
573        assert_eq!(expired_count, 1);
574
575        // Verify session is removed
576        let session = manager.get_session(&session_id).await;
577        assert!(session.is_none());
578    }
579}
ultrafast_mcp_sequential_thinking/session/mod.rs

ultrafast_mcp_sequential_thinking/session/
mod.rs
