mcp: abstract transport session storage

apps/mcp-server/src/session.rs

@@ -1,68 +1,181 @@
 use std::{collections::HashMap, sync::Arc};
 
+use async_trait::async_trait;
+use thiserror::Error;
+use time::OffsetDateTime;
 use tokio::sync::RwLock;
 use uuid::Uuid;
 
-#[derive(Clone)]
-pub struct SessionStore {
-    inner: Arc<RwLock<HashMap<String, SessionState>>>,
-}
-
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SessionState {
+    pub id: String,
     pub protocol_version: String,
     pub initialized: bool,
     pub workspace_slug: String,
     pub agent_slug: String,
+    pub created_at: OffsetDateTime,
+    pub updated_at: OffsetDateTime,
+    pub expires_at: Option<OffsetDateTime>,
 }
 
-impl SessionStore {
-    pub fn new() -> Self {
-        Self {
-            inner: Arc::new(RwLock::new(HashMap::new())),
-        }
-    }
+#[derive(Debug, Error)]
+#[error("transport session store is unavailable: {details}")]
+pub struct SessionStoreError {
+    pub details: String,
+}
 
-    pub async fn create(
+#[async_trait]
+pub trait TransportSessionStore: Send + Sync {
+    async fn create(
         &self,
         protocol_version: &str,
         workspace_slug: &str,
         agent_slug: &str,
-    ) -> String {
+        now: OffsetDateTime,
+    ) -> Result<String, SessionStoreError>;
+
+    async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError>;
+
+    async fn mark_initialized(
+        &self,
+        session_id: &str,
+        now: OffsetDateTime,
+    ) -> Result<bool, SessionStoreError>;
+
+    async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
+}
+
+pub type SharedSessionStore = Arc<dyn TransportSessionStore>;
+
+#[derive(Clone, Default)]
+pub struct InMemorySessionStore {
+    inner: Arc<RwLock<HashMap<String, SessionState>>>,
+}
+
+impl InMemorySessionStore {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+#[async_trait]
+impl TransportSessionStore for InMemorySessionStore {
+    async fn create(
+        &self,
+        protocol_version: &str,
+        workspace_slug: &str,
+        agent_slug: &str,
+        now: OffsetDateTime,
+    ) -> Result<String, SessionStoreError> {
         let session_id = Uuid::now_v7().to_string();
         let mut guard = self.inner.write().await;
 
         guard.insert(
             session_id.clone(),
             SessionState {
+                id: session_id.clone(),
                 protocol_version: protocol_version.to_owned(),
                 initialized: false,
                 workspace_slug: workspace_slug.to_owned(),
                 agent_slug: agent_slug.to_owned(),
+                created_at: now,
+                updated_at: now,
+                expires_at: None,
             },
         );
 
-        session_id
+        Ok(session_id)
     }
 
-    pub async fn get(&self, session_id: &str) -> Option<SessionState> {
+    async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
         let guard = self.inner.read().await;
-        guard.get(session_id).cloned()
+        Ok(guard.get(session_id).cloned())
     }
 
-    pub async fn mark_initialized(&self, session_id: &str) -> bool {
+    async fn mark_initialized(
+        &self,
+        session_id: &str,
+        now: OffsetDateTime,
+    ) -> Result<bool, SessionStoreError> {
         let mut guard = self.inner.write().await;
 
         if let Some(session) = guard.get_mut(session_id) {
             session.initialized = true;
-            return true;
+            session.updated_at = now;
+            return Ok(true);
         }
 
-        false
+        Ok(false)
     }
 
-    pub async fn delete(&self, session_id: &str) -> bool {
+    async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
         let mut guard = self.inner.write().await;
-        guard.remove(session_id).is_some()
+        Ok(guard.remove(session_id).is_some())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use time::format_description::well_known::Rfc3339;
+
+    use super::{InMemorySessionStore, SessionStoreError, TransportSessionStore};
+
+    fn timestamp(value: &str) -> time::OffsetDateTime {
+        time::OffsetDateTime::parse(value, &Rfc3339).unwrap()
+    }
+
+    #[tokio::test]
+    async fn creates_and_reads_transport_sessions() {
+        let store = InMemorySessionStore::new();
+        let created_at = timestamp("2026-05-01T10:00:00Z");
+
+        let session_id = store
+            .create("2025-11-25", "default", "sales", created_at)
+            .await
+            .unwrap();
+        let session = store.get(&session_id).await.unwrap().unwrap();
+
+        assert_eq!(session.id, session_id);
+        assert_eq!(session.protocol_version, "2025-11-25");
+        assert_eq!(session.workspace_slug, "default");
+        assert_eq!(session.agent_slug, "sales");
+        assert!(!session.initialized);
+        assert_eq!(session.created_at, created_at);
+        assert_eq!(session.updated_at, created_at);
+    }
+
+    #[tokio::test]
+    async fn marks_transport_sessions_initialized() {
+        let store = InMemorySessionStore::new();
+        let created_at = timestamp("2026-05-01T10:00:00Z");
+        let initialized_at = timestamp("2026-05-01T10:00:05Z");
+
+        let session_id = store
+            .create("2025-11-25", "default", "sales", created_at)
+            .await
+            .unwrap();
+
+        assert!(
+            store
+                .mark_initialized(&session_id, initialized_at)
+                .await
+                .unwrap()
+        );
+
+        let session = store.get(&session_id).await.unwrap().unwrap();
+        assert!(session.initialized);
+        assert_eq!(session.updated_at, initialized_at);
+    }
+
+    #[test]
+    fn formats_transport_session_store_error() {
+        let error = SessionStoreError {
+            details: "postgres is unavailable".to_owned(),
+        };
+
+        assert_eq!(
+            error.to_string(),
+            "transport session store is unavailable: postgres is unavailable"
+        );
     }
 }