Skip to main content

poet/mcp/
session_storage.rs

1use actix_web::Result;
2use dashmap::DashMap;
3
4use crate::mcp::session::Session;
5
6#[derive(Clone, Default)]
7pub struct SessionStorage {
8    pub sessions: DashMap<String, Session>,
9}
10
11impl SessionStorage {
12    pub async fn read(&self, session_id: &str) -> Result<Option<Session>> {
13        Ok(self.sessions.get(session_id).map(|session| session.clone()))
14    }
15
16    pub async fn store_new_session(&self, session: Session) -> Result<()> {
17        self.update_session(session).await
18    }
19
20    pub async fn terminate_session(&self, session: Session) -> Result<()> {
21        self.sessions.remove(&session.id());
22        session.terminate().await;
23
24        Ok(())
25    }
26
27    pub async fn update_session(&self, session: Session) -> Result<()> {
28        self.sessions.insert(session.id(), session);
29
30        Ok(())
31    }
32}