mcp: abstract transport session storage

TASKS.md

@@ -2,19 +2,19 @@
 
 ## Current
 
-### `feat/runtime-rate-limiting-and-backpressure`
+### `feat/distributed-mcp-session-store`
 
 Status: in_progress
 
 DoD:
-- runtime rejects overload deterministically with stable error codes
+- transport sessions are stored behind a shared store abstraction
-- unary, window, session, and async-job execution paths are concurrency-bounded
+- mcp-server no longer depends directly on process-local session storage implementation
-- configuration is explicit and documented
+- in-memory session store remains available for tests and local fallback
-- admin-api and mcp-server use the configured runtime limits
+- follow-up slices can add a Postgres-backed transport session store without changing transport handlers
 
 ## Next
 
-- `feat/distributed-mcp-session-store`
+- `feat/runtime-rate-limiting-and-backpressure`
 
 ## Backlog
 

apps/mcp-server/Cargo.toml

@@ -6,6 +6,7 @@ rust-version.workspace = true
 version.workspace = true
 
 [dependencies]
+async-trait = "0.1"
 axum.workspace = true
 base64.workspace = true
 crank-core = { path = "../../crates/crank-core" }
@@ -17,6 +18,7 @@ serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 sqlx.workspace = true
+thiserror.workspace = true
 time.workspace = true
 tokio = { workspace = true, features = ["sync"] }
 tracing.workspace = true

apps/mcp-server/src/app.rs

@@ -47,7 +47,7 @@ use crate::{
         is_response, jsonrpc_error, jsonrpc_result, method_name, negotiated_protocol_version,
         params, request_id,
     },
-    session::{SessionState, SessionStore},
+    session::{InMemorySessionStore, SessionState, SharedSessionStore},
 };
 
 const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
@@ -68,7 +68,7 @@ pub struct AppState {
     runtime: RuntimeExecutor,
     api_rate_limiter: RequestRateLimiter,
     secret_crypto: SecretCrypto,
-    sessions: SessionStore,
+    sessions: SharedSessionStore,
     allowed_origins: AllowedOrigins,
 }
 
@@ -131,6 +131,26 @@ pub fn build_app(
     secret_crypto: SecretCrypto,
     runtime: RuntimeExecutor,
     api_rate_limiter: RequestRateLimiter,
+) -> Router {
+    build_app_with_session_store(
+        registry,
+        refresh_interval,
+        public_base_url,
+        secret_crypto,
+        runtime,
+        api_rate_limiter,
+        Arc::new(InMemorySessionStore::new()),
+    )
+}
+
+pub fn build_app_with_session_store(
+    registry: PostgresRegistry,
+    refresh_interval: Duration,
+    public_base_url: Option<String>,
+    secret_crypto: SecretCrypto,
+    runtime: RuntimeExecutor,
+    api_rate_limiter: RequestRateLimiter,
+    sessions: SharedSessionStore,
 ) -> Router {
     let state = Arc::new(AppState {
         registry: registry.clone(),
@@ -138,7 +158,7 @@ pub fn build_app(
         runtime,
         api_rate_limiter,
         secret_crypto,
-        sessions: SessionStore::new(),
+        sessions,
         allowed_origins: AllowedOrigins::new(public_base_url),
     });
 
@@ -183,7 +203,10 @@ async fn mcp_get(
         Err(status) => return status.into_response(),
     };
 
-    let Some(session) = state.sessions.get(&session_id).await else {
+    let Some(session) = (match state.sessions.get(&session_id).await {
+        Ok(session) => session,
+        Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+    }) else {
         return StatusCode::NOT_FOUND.into_response();
     };
 
@@ -220,18 +243,19 @@ async fn mcp_delete(
 
     match session_id_from_headers(&headers) {
         Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
-            Some(session)
+            Ok(Some(session))
                 if session.workspace_slug == path.workspace_slug
                     && session.agent_slug == path.agent_slug =>
             {
-                if state.sessions.delete(&session_id).await {
+                match state.sessions.delete(&session_id).await {
-                    StatusCode::NO_CONTENT.into_response()
+                    Ok(true) => StatusCode::NO_CONTENT.into_response(),
-                } else {
+                    Ok(false) => StatusCode::NOT_FOUND.into_response(),
-                    StatusCode::NOT_FOUND.into_response()
+                    Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
                 }
             }
-            Some(_) => StatusCode::NOT_FOUND.into_response(),
+            Ok(Some(_)) => StatusCode::NOT_FOUND.into_response(),
-            None => StatusCode::NOT_FOUND.into_response(),
+            Ok(None) => StatusCode::NOT_FOUND.into_response(),
+            Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
         },
         Ok(None) => StatusCode::BAD_REQUEST.into_response(),
         Err(status) => status.into_response(),
@@ -284,7 +308,16 @@ async fn mcp_post(
 
     if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) {
         if let Ok(session_id) = session_id.to_str() {
-            if let Some(session) = state.sessions.get(session_id).await {
+            let session = match state.sessions.get(session_id).await {
+                Ok(session) => session,
+                Err(_) => {
+                    return with_request_id_header(
+                        StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+                        &transport_request_id,
+                    );
+                }
+            };
+            if let Some(session) = session {
                 if let Err(status) =
                     validate_session_protocol_version(&headers, &session.protocol_version)
                 {
@@ -1667,10 +1700,20 @@ async fn handle_initialize(
             Some(DEFAULT_PROTOCOL_VERSION),
         );
     };
-    let session_id = state
+    let now = OffsetDateTime::now_utc();
+    let session_id = match state
         .sessions
-        .create(protocol_version, &path.workspace_slug, &path.agent_slug)
+        .create(
-        .await;
+            protocol_version,
+            &path.workspace_slug,
+            &path.agent_slug,
+            now,
+        )
+        .await
+    {
+        Ok(session_id) => session_id,
+        Err(error) => return internal_jsonrpc_error(message, error),
+    };
 
     transport_response(
         StatusCode::OK,
@@ -1702,18 +1745,23 @@ async fn handle_initialized_notification(
 ) -> Response {
     match session_id_from_headers(headers) {
         Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
-            Some(session)
+            Ok(Some(session))
                 if session.workspace_slug == path.workspace_slug
                     && session.agent_slug == path.agent_slug =>
             {
-                if state.sessions.mark_initialized(&session_id).await {
+                match state
-                    StatusCode::ACCEPTED.into_response()
+                    .sessions
-                } else {
+                    .mark_initialized(&session_id, OffsetDateTime::now_utc())
-                    StatusCode::NOT_FOUND.into_response()
+                    .await
+                {
+                    Ok(true) => StatusCode::ACCEPTED.into_response(),
+                    Ok(false) => StatusCode::NOT_FOUND.into_response(),
+                    Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
                 }
             }
-            Some(_) => StatusCode::NOT_FOUND.into_response(),
+            Ok(Some(_)) => StatusCode::NOT_FOUND.into_response(),
-            None => StatusCode::NOT_FOUND.into_response(),
+            Ok(None) => StatusCode::NOT_FOUND.into_response(),
+            Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
         },
         Ok(None) => StatusCode::BAD_REQUEST.into_response(),
         Err(status) => status.into_response(),
@@ -1732,7 +1780,10 @@ async fn require_initialized_session(
         Err(status) => return Err(status.into_response()),
     };
 
-    let Some(session) = state.sessions.get(&session_id).await else {
+    let Some(session) = (match state.sessions.get(&session_id).await {
+        Ok(session) => session,
+        Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR.into_response()),
+    }) else {
         return Err(StatusCode::NOT_FOUND.into_response());
     };
 

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)]
+#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct SessionStore {
-    inner: Arc<RwLock<HashMap<String, SessionState>>>,
-}
-
-#[derive(Clone)]
 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 {
+#[derive(Debug, Error)]
-    pub fn new() -> Self {
+#[error("transport session store is unavailable: {details}")]
-        Self {
+pub struct SessionStoreError {
-            inner: Arc::new(RwLock::new(HashMap::new())),
+    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"
+        );
     }
 }