bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mcp: abstract transport session storage

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 17:24:59 +00:00
parent 5f68fcbd04
commit 4fa4bded7c
4 changed files with 217 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TASKS.md
+6 -6
View File
@@ -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
apps/mcp-server/Cargo.toml
+2
View File
@@ -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
+74 -23
View File
@@ -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 {
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());
};
apps/mcp-server/src/session.rs
+135 -22
View File
@@ -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"
);
}
}