Skip to main content

barbacane_lib/mcp/
session.rs

1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use parking_lot::Mutex;
5use uuid::Uuid;
6
7/// MCP session state.
8struct McpSession {
9    _client_info: Option<serde_json::Value>,
10    last_active: Instant,
11}
12
13/// Thread-safe MCP session store with TTL eviction.
14pub struct SessionStore {
15    sessions: Mutex<HashMap<String, McpSession>>,
16    ttl: Duration,
17}
18
19impl SessionStore {
20    pub fn new(ttl: Duration) -> Self {
21        Self {
22            sessions: Mutex::new(HashMap::new()),
23            ttl,
24        }
25    }
26
27    /// Create a new session and return its ID.
28    pub fn create(&self, client_info: Option<serde_json::Value>) -> String {
29        let id = Uuid::new_v4().to_string();
30        let session = McpSession {
31            _client_info: client_info,
32            last_active: Instant::now(),
33        };
34        self.sessions.lock().insert(id.clone(), session);
35        id
36    }
37
38    /// Touch a session (update last_active). Returns false if session doesn't exist.
39    pub fn touch(&self, id: &str) -> bool {
40        let mut sessions = self.sessions.lock();
41        if let Some(session) = sessions.get_mut(id) {
42            session.last_active = Instant::now();
43            true
44        } else {
45            false
46        }
47    }
48
49    /// Remove a session.
50    pub fn remove(&self, id: &str) {
51        self.sessions.lock().remove(id);
52    }
53
54    /// Evict expired sessions. Call periodically from a background task.
55    pub fn evict_expired(&self) {
56        let now = Instant::now();
57        self.sessions
58            .lock()
59            .retain(|_, session| now.duration_since(session.last_active) < self.ttl);
60    }
61
62    /// Number of active sessions.
63    #[cfg(test)]
64    pub fn len(&self) -> usize {
65        self.sessions.lock().len()
66    }
67
68    /// Whether the session store is empty.
69    #[cfg(test)]
70    pub fn is_empty(&self) -> bool {
71        self.sessions.lock().is_empty()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn create_and_touch() {
81        let store = SessionStore::new(Duration::from_secs(60));
82        let id = store.create(None);
83        assert!(store.touch(&id));
84        assert!(!store.touch("nonexistent"));
85        assert_eq!(store.len(), 1);
86    }
87
88    #[test]
89    fn remove_session() {
90        let store = SessionStore::new(Duration::from_secs(60));
91        let id = store.create(None);
92        assert_eq!(store.len(), 1);
93        store.remove(&id);
94        assert_eq!(store.len(), 0);
95        assert!(!store.touch(&id));
96    }
97
98    #[test]
99    fn evict_expired() {
100        let store = SessionStore::new(Duration::from_millis(1));
101        store.create(None);
102        store.create(None);
103        assert_eq!(store.len(), 2);
104        std::thread::sleep(Duration::from_millis(10));
105        store.evict_expired();
106        assert_eq!(store.len(), 0);
107    }
108}