webterm_agent/models/
session_registry.rs

1use crate::models::agent_error::AgentError;
2use crate::models::session::Session;
3use std::collections::HashMap;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, OnceLock};
6use tokio::sync::{Mutex, RwLock};
7use tracing::debug;
8use webterm_core::types::SessionId;
9
10// Start from 1. Frontend may send "0" session ID to request a new session
11static NEXT_SESSION_ID: AtomicU64 = AtomicU64::new(1);
12
13pub struct SessionRegistry {
14    map: RwLock<HashMap<SessionId, Arc<Mutex<Session>>>>,
15}
16
17impl SessionRegistry {
18    pub(crate) async fn singleton() -> &'static Self {
19        static INSTANCE: OnceLock<SessionRegistry> = OnceLock::new();
20        INSTANCE.get_or_init(|| Self {
21            map: RwLock::new(HashMap::new()),
22        })
23    }
24
25    pub async fn find(session_id: SessionId) -> Result<Arc<Mutex<Session>>, AgentError> {
26        let registry = Self::singleton().await;
27        registry
28            .map
29            .read()
30            .await
31            .get(&session_id)
32            .ok_or(AgentError::SessionNotFound(Some(session_id)))
33            .cloned()
34    }
35
36    pub async fn build_session() -> Result<Arc<Mutex<Session>>, AgentError> {
37        let registry = Self::singleton().await;
38        let session = Arc::new(Mutex::new(Session::new()));
39        let session_ = session.lock().await;
40        debug!("Registered session {:?}", session_.session_id());
41        registry
42            .map
43            .write()
44            .await
45            .insert(session_.session_id(), session.clone());
46
47        Ok(session.clone())
48    }
49
50    pub fn next_session_id() -> SessionId {
51        SessionId(NEXT_SESSION_ID.fetch_add(1, Ordering::SeqCst))
52    }
53}