jevsnes.git / packages / mcp / src / sessions.rs
sessions.rsannotatedsessions.rssource95 lines · 3.7 KB · raw

MCP sessions that outlive the window.

A client is told about a changed tool list over a stream it holds open, and a stream belongs to a session, so the server keeps sessions. Left at that, a restarted app would not know any of its clients, and each would have to notice and start again. rmcp restores a session it does not know from a [SessionStore] if it is given one, so this is one: a file of who has introduced themselves, under $XDG_STATE_HOME/jev.

10use std::collections::BTreeMap;
11use std::path::PathBuf;
12use std::sync::{Mutex, PoisonError};
13use std::time::{SystemTime, UNIX_EPOCH};
14use std::{env, fs, io};

16use rmcp::transport::streamable_http_server::session::{
17    SessionState, SessionStore, SessionStoreError,
18};

How many sessions are remembered. A client that goes away without saying so leaves its entry behind, so the file is kept to the most recent few rather than left to grow; one window has a handful of clients at most.

23const REMEMBERED: usize = 32;

25#[derive(serde::Serialize, serde::Deserialize)]
26struct Remembered {

Seconds since the epoch, to tell which are the most recent.

28    stored: u64,
29    state: SessionState,
30}

32pub struct FileSessionStore {
33    path: PathBuf,
34    sessions: Mutex<BTreeMap<String, Remembered>>,
35}
36
37impl FileSessionStore {

Open the store, reading what an earlier window left. A file that cannot be read is a store with nobody in it: those clients introduce themselves again, which is all that would have been saved.

41    pub fn open() -> Self {
42        let state_home = env::var_os("XDG_STATE_HOME")
43            .map(PathBuf::from)
44            .filter(|path| path.is_absolute())
45            .or_else(|| Some(PathBuf::from(env::var_os("HOME")?).join(".local/state")))
46            .unwrap_or_else(env::temp_dir);
47        let path = state_home.join("jev/mcp-sessions.json");
48        let sessions = fs::read(&path)
49            .ok()
50            .and_then(|bytes| serde_json::from_slice(&bytes).ok())
51            .unwrap_or_default();
52        Self { path, sessions: Mutex::new(sessions) }
53    }

55    fn change(&self, change: impl FnOnce(&mut BTreeMap<String, Remembered>)) -> io::Result<()> {
56        let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
57        change(&mut sessions);
58        while sessions.len() > REMEMBERED {
59            let oldest = sessions
60                .iter()
61                .min_by_key(|(_, remembered)| remembered.stored)
62                .map(|(id, _)| id.clone())
63                .expect("more than REMEMBERED entries is at least one");
64            sessions.remove(&oldest);
65        }
66        if let Some(dir) = self.path.parent() {
67            fs::create_dir_all(dir)?;
68        }
69        let partial = self.path.with_extension("partial");
70        fs::write(&partial, serde_json::to_vec(&*sessions)?)?;
71        fs::rename(partial, &self.path)
72    }
73}
74
75#[async_trait::async_trait]
76impl SessionStore for FileSessionStore {
77    async fn load(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
78        let sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
79        Ok(sessions.get(session_id).map(|remembered| remembered.state.clone()))
80    }
81
82    async fn store(&self, session_id: &str, state: &SessionState) -> Result<(), SessionStoreError> {
83        let stored = SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |d| d.as_secs());
84        let remembered = Remembered { stored, state: state.clone() };
85        Ok(self.change(|sessions| {
86            sessions.insert(session_id.to_owned(), remembered);
87        })?)
88    }
89
90    async fn delete(&self, session_id: &str) -> Result<(), SessionStoreError> {
91        Ok(self.change(|sessions| {
92            sessions.remove(session_id);
93        })?)
94    }
95}