1use crate::eval::{EvalContext, EvalError};
9use crate::protocol::{ReplMode, SessionId};
10use std::collections::HashMap;
11use std::sync::{Arc, RwLock};
12use thiserror::Error;
13
14#[derive(Debug, Error)]
16pub enum SessionError {
17 #[error("Session not found: {0}")]
18 NotFound(String),
19
20 #[error("Session already exists: {0}")]
21 AlreadyExists(String),
22
23 #[error("Lock poisoned")]
24 LockPoisoned,
25
26 #[error("Evaluation failed")]
27 EvalFailed(#[from] EvalError),
28}
29
30pub type Result<T> = std::result::Result<T, SessionError>;
31
32#[derive(Debug, Clone)]
34pub struct SessionInfo {
35 pub id: SessionId,
37
38 pub name: Option<String>,
40
41 pub mode: ReplMode,
43
44 pub eval_count: u64,
46
47 pub created_at: u64,
49
50 pub last_active_at: u64,
52
53 pub timeout_ms: u64,
55}
56
57#[derive(Debug, Clone)]
59struct SessionMetadata {
60 name: Option<String>,
62 created_at: u64,
64 last_active_at: u64,
66 timeout_ms: u64,
68}
69
70impl Default for SessionMetadata {
71 fn default() -> Self {
72 Self {
73 name: None,
74 created_at: std::time::SystemTime::now()
75 .duration_since(std::time::UNIX_EPOCH)
76 .unwrap()
77 .as_millis() as u64,
78 last_active_at: std::time::SystemTime::now()
79 .duration_since(std::time::UNIX_EPOCH)
80 .unwrap()
81 .as_millis() as u64,
82 timeout_ms: 3_600_000, }
84 }
85}
86
87#[derive(Debug, Clone)]
92pub struct SessionManager {
93 sessions: Arc<RwLock<HashMap<SessionId, EvalContext>>>,
95 metadata: Arc<RwLock<HashMap<SessionId, SessionMetadata>>>,
97}
98
99impl SessionManager {
100 pub fn new() -> Self {
102 Self {
103 sessions: Arc::new(RwLock::new(HashMap::new())),
104 metadata: Arc::new(RwLock::new(HashMap::new())),
105 }
106 }
107
108 pub fn create(&self, session_id: SessionId, mode: ReplMode) -> Result<SessionId> {
125 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
126
127 if sessions.contains_key(&session_id) {
128 return Err(SessionError::AlreadyExists(session_id.to_string()));
129 }
130
131 let context = EvalContext::new(session_id.clone(), mode);
134 sessions.insert(session_id.clone(), context);
135
136 let mut metadata = self.metadata.write().map_err(|_| SessionError::LockPoisoned)?;
138 metadata.insert(session_id.clone(), SessionMetadata::default());
139
140 Ok(session_id)
141 }
142
143 pub fn create_with_compilation(
153 &self,
154 session_id: SessionId,
155 mode: ReplMode,
156 ) -> Result<SessionId> {
157 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
158
159 if sessions.contains_key(&session_id) {
160 return Err(SessionError::AlreadyExists(session_id.to_string()));
161 }
162
163 let context = EvalContext::with_compilation(session_id.clone(), mode)?;
165 sessions.insert(session_id.clone(), context);
166
167 let mut metadata = self.metadata.write().map_err(|_| SessionError::LockPoisoned)?;
169 metadata.insert(session_id.clone(), SessionMetadata::default());
170
171 Ok(session_id)
172 }
173
174 pub async fn eval(
191 &self,
192 session_id: &SessionId,
193 code: impl AsRef<str>,
194 ) -> Result<crate::eval::EvalResult> {
195 let mut context = {
197 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
198
199 sessions
200 .get(session_id)
201 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?
202 .clone()
203 };
204
205 let result = context.eval(code.as_ref()).await?;
207
208 {
210 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
211
212 sessions.insert(session_id.clone(), context.clone());
213 }
214
215 {
217 let mut metadata = self.metadata.write().map_err(|_| SessionError::LockPoisoned)?;
218 if let Some(meta) = metadata.get_mut(session_id) {
219 meta.last_active_at = std::time::SystemTime::now()
220 .duration_since(std::time::UNIX_EPOCH)
221 .unwrap()
222 .as_millis() as u64;
223 }
224 }
225
226 Ok(result)
227 }
228
229 pub fn clone_session(&self, source_id: &SessionId, target_id: SessionId) -> Result<SessionId> {
239 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
240
241 let source_context =
242 sessions.get(source_id).ok_or_else(|| SessionError::NotFound(source_id.to_string()))?;
243
244 if sessions.contains_key(&target_id) {
245 return Err(SessionError::AlreadyExists(target_id.to_string()));
246 }
247
248 let cloned_context = source_context.clone_to(target_id.clone());
249 sessions.insert(target_id.clone(), cloned_context);
250
251 Ok(target_id)
252 }
253
254 pub fn close(&self, session_id: &SessionId) -> Result<()> {
262 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
263
264 sessions
265 .remove(session_id)
266 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
267
268 Ok(())
269 }
270
271 pub fn list(&self) -> Result<Vec<SessionInfo>> {
279 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
280 let metadata = self.metadata.read().map_err(|_| SessionError::LockPoisoned)?;
281
282 let mut infos: Vec<SessionInfo> = sessions
283 .iter()
284 .map(|(id, ctx)| {
285 let (tier1, tier2, _) = ctx.stats();
286 let meta = metadata.get(id).cloned().unwrap_or_default();
287
288 SessionInfo {
289 id: id.clone(),
290 name: meta.name,
291 mode: ctx.mode(),
292 eval_count: tier1 + tier2,
293 created_at: meta.created_at,
294 last_active_at: meta.last_active_at,
295 timeout_ms: meta.timeout_ms,
296 }
297 })
298 .collect();
299
300 infos.sort_by(|a, b| a.id.cmp(&b.id));
302
303 Ok(infos)
304 }
305
306 pub fn get_info(&self, session_id: &SessionId) -> Result<SessionInfo> {
312 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
313 let metadata = self.metadata.read().map_err(|_| SessionError::LockPoisoned)?;
314
315 let ctx = sessions
316 .get(session_id)
317 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
318
319 let (tier1, tier2, _) = ctx.stats();
320 let meta = metadata.get(session_id).cloned().unwrap_or_default();
321
322 Ok(SessionInfo {
323 id: session_id.clone(),
324 name: meta.name,
325 mode: ctx.mode(),
326 eval_count: tier1 + tier2,
327 created_at: meta.created_at,
328 last_active_at: meta.last_active_at,
329 timeout_ms: meta.timeout_ms,
330 })
331 }
332
333 pub fn get_stats_collector(
342 &self,
343 session_id: &SessionId,
344 ) -> Result<Arc<std::sync::Mutex<crate::eval::EvalMetrics>>> {
345 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
346
347 let ctx = sessions
348 .get(session_id)
349 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
350
351 Ok(ctx.stats_collector())
352 }
353
354 pub fn get_usage_metrics(
363 &self,
364 session_id: &SessionId,
365 ) -> Result<Arc<std::sync::Mutex<crate::metrics::UsageMetrics>>> {
366 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
367
368 let ctx = sessions
369 .get(session_id)
370 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
371
372 Ok(ctx.usage_metrics())
373 }
374
375 pub fn get_resource_stats(
384 &self,
385 session_id: &SessionId,
386 ) -> Result<(Option<crate::session::DirStats>, Option<crate::cache::CacheStats>)> {
387 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
388
389 let ctx = sessions
390 .get(session_id)
391 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
392
393 Ok((ctx.session_dir_stats(), ctx.artifact_cache_stats()))
394 }
395
396 pub fn get_subprocess_stats(
405 &self,
406 session_id: &SessionId,
407 ) -> Result<Option<crate::metrics::SubprocessMetricsSnapshot>> {
408 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
409
410 let ctx = sessions
411 .get(session_id)
412 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
413
414 Ok(ctx.subprocess_stats())
415 }
416
417 pub fn count(&self) -> Result<usize> {
423 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
424
425 Ok(sessions.len())
426 }
427
428 pub fn exists(&self, session_id: &SessionId) -> Result<bool> {
434 let sessions = self.sessions.read().map_err(|_| SessionError::LockPoisoned)?;
435
436 Ok(sessions.contains_key(session_id))
437 }
438
439 pub fn close_all(&self) -> Result<usize> {
447 let mut sessions = self.sessions.write().map_err(|_| SessionError::LockPoisoned)?;
448
449 let count = sessions.len();
450 sessions.clear();
451
452 Ok(count)
453 }
454}
455
456impl Default for SessionManager {
457 fn default() -> Self {
458 Self::new()
459 }
460}
461
462#[cfg(test)]
463mod tests {
464 use super::*;
465
466 #[test]
467 fn test_create_session() {
468 let manager = SessionManager::new();
469 let id = manager.create(SessionId::new("test-1"), ReplMode::Lisp).unwrap();
470
471 assert_eq!(id, SessionId::new("test-1"));
472 assert!(manager.exists(&id).unwrap());
473 }
474
475 #[test]
476 fn test_create_duplicate_session() {
477 let manager = SessionManager::new();
478 manager.create(SessionId::new("test-1"), ReplMode::Lisp).unwrap();
479
480 let result = manager.create(SessionId::new("test-1"), ReplMode::Lisp);
481 assert!(matches!(result, Err(SessionError::AlreadyExists(_))));
482 }
483
484 #[tokio::test]
485 async fn test_eval() {
486 let manager = SessionManager::new();
487 manager.create(SessionId::new("test"), ReplMode::Lisp).unwrap();
488
489 let result = manager.eval(&SessionId::new("test"), "(+ 1 2)").await.unwrap();
490
491 assert_eq!(result.value, "3");
492 }
493
494 #[tokio::test]
495 async fn test_eval_not_found() {
496 let manager = SessionManager::new();
497
498 let result = manager.eval(&SessionId::new("nonexistent"), "test").await;
499
500 assert!(matches!(result, Err(SessionError::NotFound(_))));
501 }
502
503 #[test]
504 fn test_clone_session() {
505 let manager = SessionManager::new();
506 manager.create(SessionId::new("source"), ReplMode::Lisp).unwrap();
507
508 let cloned_id =
509 manager.clone_session(&SessionId::new("source"), SessionId::new("target")).unwrap();
510
511 assert_eq!(cloned_id, SessionId::new("target"));
512 assert!(manager.exists(&SessionId::new("target")).unwrap());
513 }
514
515 #[test]
516 fn test_clone_nonexistent_session() {
517 let manager = SessionManager::new();
518
519 let result =
520 manager.clone_session(&SessionId::new("nonexistent"), SessionId::new("target"));
521
522 assert!(matches!(result, Err(SessionError::NotFound(_))));
523 }
524
525 #[test]
526 fn test_clone_to_existing_session() {
527 let manager = SessionManager::new();
528 manager.create(SessionId::new("source"), ReplMode::Lisp).unwrap();
529 manager.create(SessionId::new("target"), ReplMode::Lisp).unwrap();
530
531 let result = manager.clone_session(&SessionId::new("source"), SessionId::new("target"));
532
533 assert!(matches!(result, Err(SessionError::AlreadyExists(_))));
534 }
535
536 #[test]
537 fn test_close_session() {
538 let manager = SessionManager::new();
539 manager.create(SessionId::new("test"), ReplMode::Lisp).unwrap();
540
541 assert!(manager.exists(&SessionId::new("test")).unwrap());
542
543 manager.close(&SessionId::new("test")).unwrap();
544
545 assert!(!manager.exists(&SessionId::new("test")).unwrap());
546 }
547
548 #[test]
549 fn test_close_nonexistent_session() {
550 let manager = SessionManager::new();
551
552 let result = manager.close(&SessionId::new("nonexistent"));
553
554 assert!(matches!(result, Err(SessionError::NotFound(_))));
555 }
556
557 #[test]
558 fn test_list_sessions() {
559 let manager = SessionManager::new();
560 manager.create(SessionId::new("session-1"), ReplMode::Lisp).unwrap();
561 manager.create(SessionId::new("session-2"), ReplMode::Sexpr).unwrap();
562
563 let sessions = manager.list().unwrap();
564
565 assert_eq!(sessions.len(), 2);
566 assert_eq!(sessions[0].id, SessionId::new("session-1"));
567 assert_eq!(sessions[0].mode, ReplMode::Lisp);
568 assert_eq!(sessions[1].id, SessionId::new("session-2"));
569 assert_eq!(sessions[1].mode, ReplMode::Sexpr);
570 }
571
572 #[test]
573 fn test_list_empty() {
574 let manager = SessionManager::new();
575 let sessions = manager.list().unwrap();
576
577 assert_eq!(sessions.len(), 0);
578 }
579
580 #[test]
581 fn test_get_info() {
582 let manager = SessionManager::new();
583 manager.create(SessionId::new("test"), ReplMode::Lisp).unwrap();
584
585 let info = manager.get_info(&SessionId::new("test")).unwrap();
586
587 assert_eq!(info.id, SessionId::new("test"));
588 assert_eq!(info.mode, ReplMode::Lisp);
589 assert_eq!(info.eval_count, 0);
590 }
591
592 #[test]
593 fn test_get_info_nonexistent() {
594 let manager = SessionManager::new();
595
596 let result = manager.get_info(&SessionId::new("nonexistent"));
597
598 assert!(matches!(result, Err(SessionError::NotFound(_))));
599 }
600
601 #[tokio::test]
602 async fn test_eval_count_tracking() {
603 let manager = SessionManager::new();
604 manager.create(SessionId::new("test"), ReplMode::Lisp).unwrap();
605
606 manager.eval(&SessionId::new("test"), "(+ 1 2)").await.unwrap();
608
609 manager.eval(&SessionId::new("test"), "(* 3 4)").await.unwrap();
610
611 let info = manager.get_info(&SessionId::new("test")).unwrap();
612 assert_eq!(info.eval_count, 2);
613 }
614
615 #[test]
616 fn test_count() {
617 let manager = SessionManager::new();
618
619 assert_eq!(manager.count().unwrap(), 0);
620
621 manager.create(SessionId::new("session-1"), ReplMode::Lisp).unwrap();
622 assert_eq!(manager.count().unwrap(), 1);
623
624 manager.create(SessionId::new("session-2"), ReplMode::Sexpr).unwrap();
625 assert_eq!(manager.count().unwrap(), 2);
626
627 manager.close(&SessionId::new("session-1")).unwrap();
628 assert_eq!(manager.count().unwrap(), 1);
629 }
630
631 #[test]
632 fn test_close_all() {
633 let manager = SessionManager::new();
634 manager.create(SessionId::new("session-1"), ReplMode::Lisp).unwrap();
635 manager.create(SessionId::new("session-2"), ReplMode::Sexpr).unwrap();
636 manager.create(SessionId::new("session-3"), ReplMode::Lisp).unwrap();
637
638 assert_eq!(manager.count().unwrap(), 3);
639
640 let closed = manager.close_all().unwrap();
641
642 assert_eq!(closed, 3);
643 assert_eq!(manager.count().unwrap(), 0);
644 }
645
646 #[test]
647 fn test_thread_safety() {
648 use std::sync::Arc;
649 use std::thread;
650
651 let manager = Arc::new(SessionManager::new());
652 let mut handles = vec![];
653
654 for i in 0..10 {
656 let manager_clone = Arc::clone(&manager);
657 let handle = thread::spawn(move || {
658 manager_clone
659 .create(SessionId::new(format!("session-{}", i)), ReplMode::Lisp)
660 .unwrap();
661 });
662 handles.push(handle);
663 }
664
665 for handle in handles {
666 handle.join().unwrap();
667 }
668
669 assert_eq!(manager.count().unwrap(), 10);
670 }
671}