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
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());
};