Skip to main content

mcpr_core/protocol/
session.rs

1use std::future::Future;
2use std::sync::Arc;
3
4use chrono::{DateTime, Utc};
5use dashmap::DashMap;
6
7/// Observed MCP session state — inferred from method calls passing through the proxy.
8/// The proxy doesn't control transitions; it observes them.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub enum SessionState {
11    Created,
12    Initialized,
13    Active,
14    Closed,
15}
16
17/// Client identity extracted from the MCP `initialize` request's `clientInfo` param.
18#[derive(Debug, Clone)]
19pub struct ClientInfo {
20    pub name: String,
21    pub version: Option<String>,
22}
23
24/// Observed session metadata tracked by the proxy for debugging/observability.
25#[derive(Debug, Clone)]
26pub struct SessionInfo {
27    pub id: String,
28    pub state: SessionState,
29    pub client_info: Option<ClientInfo>,
30    pub created_at: DateTime<Utc>,
31    pub last_active: DateTime<Utc>,
32    pub request_count: u64,
33}
34
35impl SessionInfo {
36    pub fn new(id: String) -> Self {
37        let now = Utc::now();
38        Self {
39            id,
40            state: SessionState::Created,
41            client_info: None,
42            created_at: now,
43            last_active: now,
44            request_count: 0,
45        }
46    }
47}
48
49/// Trait for session storage backends.
50/// Async to support I/O-backed stores (Redis, database, logging).
51pub trait SessionStore: Send + Sync + 'static {
52    fn create(&self, id: &str) -> impl Future<Output = SessionInfo> + Send;
53    fn get(&self, id: &str) -> impl Future<Output = Option<SessionInfo>> + Send;
54    fn touch(&self, id: &str) -> impl Future<Output = ()> + Send;
55    fn update_state(&self, id: &str, state: SessionState) -> impl Future<Output = ()> + Send;
56    fn set_client_info(&self, id: &str, info: ClientInfo) -> impl Future<Output = ()> + Send;
57    fn remove(&self, id: &str) -> impl Future<Output = ()> + Send;
58    fn list(&self) -> impl Future<Output = Vec<SessionInfo>> + Send;
59}
60
61/// In-memory session store backed by DashMap for lock-free concurrent access.
62#[derive(Clone)]
63pub struct MemorySessionStore {
64    sessions: Arc<DashMap<String, SessionInfo>>,
65}
66
67impl MemorySessionStore {
68    pub fn new() -> Self {
69        Self {
70            sessions: Arc::new(DashMap::new()),
71        }
72    }
73
74    /// Sync access to session list — for use in non-async contexts (TUI rendering).
75    pub fn list_sync(&self) -> Vec<SessionInfo> {
76        self.sessions.iter().map(|r| r.value().clone()).collect()
77    }
78}
79
80impl Default for MemorySessionStore {
81    fn default() -> Self {
82        Self::new()
83    }
84}
85
86impl SessionStore for MemorySessionStore {
87    async fn create(&self, id: &str) -> SessionInfo {
88        let info = SessionInfo::new(id.to_string());
89        self.sessions.insert(id.to_string(), info.clone());
90        info
91    }
92
93    async fn get(&self, id: &str) -> Option<SessionInfo> {
94        self.sessions.get(id).map(|r| r.clone())
95    }
96
97    async fn touch(&self, id: &str) {
98        if let Some(mut entry) = self.sessions.get_mut(id) {
99            entry.last_active = Utc::now();
100            entry.request_count += 1;
101        }
102    }
103
104    async fn update_state(&self, id: &str, state: SessionState) {
105        if let Some(mut entry) = self.sessions.get_mut(id) {
106            entry.state = state;
107            entry.last_active = Utc::now();
108        }
109    }
110
111    async fn set_client_info(&self, id: &str, info: ClientInfo) {
112        if let Some(mut entry) = self.sessions.get_mut(id) {
113            entry.client_info = Some(info);
114        }
115    }
116
117    async fn remove(&self, id: &str) {
118        self.sessions.remove(id);
119    }
120
121    async fn list(&self) -> Vec<SessionInfo> {
122        self.sessions.iter().map(|r| r.value().clone()).collect()
123    }
124}
125
126/// Extract `clientInfo` from MCP initialize request params.
127pub fn parse_client_info(params: &serde_json::Value) -> Option<ClientInfo> {
128    let client_info = params.get("clientInfo")?;
129    let name = client_info.get("name")?.as_str()?.to_string();
130    let version = client_info
131        .get("version")
132        .and_then(|v| v.as_str())
133        .map(String::from);
134    Some(ClientInfo { name, version })
135}