diff --git a/TASKS.md b/TASKS.md index dbd22c5..f7c3dca 100644 --- a/TASKS.md +++ b/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 -- unary, window, session, and async-job execution paths are concurrency-bounded -- configuration is explicit and documented -- admin-api and mcp-server use the configured runtime limits +- transport sessions are stored behind a shared store abstraction +- mcp-server no longer depends directly on process-local session storage implementation +- in-memory session store remains available for tests and local fallback +- 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 diff --git a/apps/mcp-server/Cargo.toml b/apps/mcp-server/Cargo.toml index 22f2140..345aa30 100644 --- a/apps/mcp-server/Cargo.toml +++ b/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 diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 5ea21f0..7c7d0b6 100644 --- a/apps/mcp-server/src/app.rs +++ b/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, + 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 { - StatusCode::NO_CONTENT.into_response() - } else { - StatusCode::NOT_FOUND.into_response() + match state.sessions.delete(&session_id).await { + Ok(true) => StatusCode::NO_CONTENT.into_response(), + Ok(false) => StatusCode::NOT_FOUND.into_response(), + Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(), } } - Some(_) => StatusCode::NOT_FOUND.into_response(), - None => StatusCode::NOT_FOUND.into_response(), + Ok(Some(_)) => 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) - .await; + .create( + 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 { - StatusCode::ACCEPTED.into_response() - } else { - StatusCode::NOT_FOUND.into_response() + match state + .sessions + .mark_initialized(&session_id, OffsetDateTime::now_utc()) + .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(), - None => StatusCode::NOT_FOUND.into_response(), + Ok(Some(_)) => 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()); }; diff --git a/apps/mcp-server/src/session.rs b/apps/mcp-server/src/session.rs index 4f55a37..bc5f35d 100644 --- a/apps/mcp-server/src/session.rs +++ b/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>>, -} - -#[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, } -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; + + async fn get(&self, session_id: &str) -> Result, SessionStoreError>; + + async fn mark_initialized( + &self, + session_id: &str, + now: OffsetDateTime, + ) -> Result; + + async fn delete(&self, session_id: &str) -> Result; +} + +pub type SharedSessionStore = Arc; + +#[derive(Clone, Default)] +pub struct InMemorySessionStore { + inner: Arc>>, +} + +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 { 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 { + async fn get(&self, session_id: &str) -> Result, 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 { 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 { 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" + ); } }