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

mcp: evict expired transport sessions

Browse Source
This commit is contained in:
a.tolmachev
2026-05-02 08:26:59 +00:00
parent 77f83a1dc1
commit fdad20a8d1
4 changed files with 204 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TASKS.md
+2 -1
View File
@@ -10,7 +10,8 @@ DoD:
- 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
- production mcp-server uses a Postgres-backed transport session store
- expired transport sessions are evicted on read and no longer leak indefinitely
## Next
apps/mcp-server/src/app.rs
+3
View File
@@ -54,6 +54,7 @@ const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
const MAX_REQUEST_ID_LEN: usize = 128;
const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000;
#[derive(Clone, Copy)]
enum ResponseMode {
@@ -1682,6 +1683,7 @@ async fn handle_initialize(
);
};
let now = OffsetDateTime::now_utc();
let expires_at = add_millis(now, TRANSPORT_SESSION_TTL_MS);
let session_id = match state
.sessions
.create(
@@ -1689,6 +1691,7 @@ async fn handle_initialize(
&path.workspace_slug,
&path.agent_slug,
now,
Some(expires_at),
)
.await
{
apps/mcp-server/src/main.rs
+72 -2
View File
@@ -158,7 +158,10 @@ mod tests {
use tokio::time::sleep;
use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*};
use crate::{app::build_app, session::InMemorySessionStore};
use crate::{
app::build_app,
session::{InMemorySessionStore, SharedSessionStore, TransportSessionStore},
};
fn test_workspace_id() -> WorkspaceId {
WorkspaceId::new("ws_default")
@@ -222,6 +225,22 @@ mod tests {
refresh_interval: Duration,
public_base_url: Option<String>,
rate_limit_config: RequestRateLimitConfig,
) -> axum::Router {
build_test_app_with_store(
registry,
refresh_interval,
public_base_url,
rate_limit_config,
std::sync::Arc::new(InMemorySessionStore::default()),
)
}
fn build_test_app_with_store(
registry: PostgresRegistry,
refresh_interval: Duration,
public_base_url: Option<String>,
rate_limit_config: RequestRateLimitConfig,
sessions: SharedSessionStore,
) -> axum::Router {
build_app(
registry,
@@ -230,7 +249,7 @@ mod tests {
SecretCrypto::new("test-master-key").unwrap(),
RuntimeExecutor::new(),
RequestRateLimiter::new(rate_limit_config),
std::sync::Arc::new(InMemorySessionStore::default()),
sessions,
)
}
@@ -877,6 +896,57 @@ mod tests {
);
}
#[tokio::test]
async fn get_returns_not_found_for_expired_transport_session() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "sales-expired-session", vec![]).await;
let api_key = create_platform_api_key(
&registry,
"mcp-expired-session",
&[PlatformApiKeyScope::Read],
)
.await;
let session_store = Arc::new(InMemorySessionStore::default());
let session_id = session_store
.create(
"2025-11-25",
test_workspace_slug(),
"sales-expired-session",
OffsetDateTime::parse("2026-05-01T10:00:00Z", &Rfc3339).unwrap(),
Some(OffsetDateTime::parse("2026-05-01T10:00:01Z", &Rfc3339).unwrap()),
)
.await
.unwrap();
session_store
.mark_initialized(
&session_id,
OffsetDateTime::parse("2026-05-01T10:00:01Z", &Rfc3339).unwrap(),
)
.await
.unwrap();
let base_url = spawn_mcp_server(build_test_app_with_store(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
RequestRateLimitConfig::new(10_000, 10_000).unwrap(),
session_store,
))
.await;
let client = reqwest::Client::new();
let response = client
.get(agent_mcp_url(&base_url, "sales-expired-session"))
.header(header::ACCEPT, "text/event-stream")
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
.header("MCP-Session-Id", &session_id)
.header("MCP-Protocol-Version", "2025-11-25")
.send()
.await
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn get_requires_session_header() {
let registry = test_registry().await;
apps/mcp-server/src/session.rs
+127 -9
View File
@@ -38,6 +38,7 @@ pub trait TransportSessionStore: Send + Sync {
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError>;
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError>;
@@ -100,6 +101,7 @@ impl TransportSessionStore for InMemorySessionStore {
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
let mut guard = self.inner.write().await;
@@ -114,7 +116,7 @@ impl TransportSessionStore for InMemorySessionStore {
agent_slug: agent_slug.to_owned(),
created_at: now,
updated_at: now,
expires_at: None,
expires_at,
},
);
@@ -122,8 +124,20 @@ impl TransportSessionStore for InMemorySessionStore {
}
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
let guard = self.inner.read().await;
Ok(guard.get(session_id).cloned())
{
let guard = self.inner.read().await;
if let Some(session) = guard.get(session_id) {
if !is_expired(session, OffsetDateTime::now_utc()) {
return Ok(Some(session.clone()));
}
} else {
return Ok(None);
}
}
let mut guard = self.inner.write().await;
guard.remove(session_id);
Ok(None)
}
async fn mark_initialized(
@@ -156,6 +170,7 @@ impl TransportSessionStore for PostgresTransportSessionStore {
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
query(
@@ -169,7 +184,7 @@ impl TransportSessionStore for PostgresTransportSessionStore {
updated_at,
expires_at
) values (
$1, $2, false, $3, $4, $5::timestamptz, $5::timestamptz, null
$1, $2, false, $3, $4, $5::timestamptz, $5::timestamptz, $6::timestamptz
)",
)
.bind(&session_id)
@@ -177,6 +192,7 @@ impl TransportSessionStore for PostgresTransportSessionStore {
.bind(workspace_slug)
.bind(agent_slug)
.bind(now)
.bind(expires_at)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
@@ -207,7 +223,7 @@ impl TransportSessionStore for PostgresTransportSessionStore {
details: error.to_string(),
})?;
Ok(row.map(|row| SessionState {
let Some(session) = row.map(|row| SessionState {
id: row.id,
protocol_version: row.protocol_version,
initialized: row.initialized,
@@ -216,7 +232,22 @@ impl TransportSessionStore for PostgresTransportSessionStore {
created_at: row.created_at,
updated_at: row.updated_at,
expires_at: row.expires_at,
}))
}) else {
return Ok(None);
};
if is_expired(&session, OffsetDateTime::now_utc()) {
query("delete from mcp_transport_sessions where id = $1")
.bind(session_id)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
return Ok(None);
}
Ok(Some(session))
}
async fn mark_initialized(
@@ -295,6 +326,12 @@ async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreErro
Ok(())
}
fn is_expired(session: &SessionState, now: OffsetDateTime) -> bool {
session
.expires_at
.is_some_and(|expires_at| expires_at <= now)
}
#[cfg(test)]
mod tests {
use std::env;
@@ -322,7 +359,7 @@ mod tests {
let created_at = timestamp("2026-05-01T10:00:00Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at)
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
let session = store.get(&session_id).await.unwrap().unwrap();
@@ -343,7 +380,7 @@ mod tests {
let initialized_at = timestamp("2026-05-01T10:00:05Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at)
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
@@ -359,6 +396,26 @@ mod tests {
assert_eq!(session.updated_at, initialized_at);
}
#[tokio::test]
async fn drops_expired_in_memory_transport_sessions_on_read() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let expires_at = timestamp("2026-05-01T10:00:01Z");
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(expires_at),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
}
#[test]
fn formats_transport_session_store_error() {
let error = SessionStoreError {
@@ -401,7 +458,13 @@ mod tests {
let created_at = timestamp("2026-05-01T10:00:00Z");
let initialized_at = timestamp("2026-05-01T10:00:05Z");
let session_id = store_a
.create("2025-11-25", "default", "sales", created_at)
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(timestamp("2026-05-02T10:00:00Z")),
)
.await
.unwrap();
store_a
@@ -423,4 +486,59 @@ mod tests {
assert_eq!(session.workspace_slug, "default");
assert_eq!(session.agent_slug, "sales");
}
#[tokio::test]
async fn postgres_transport_sessions_evict_expired_rows_on_read() {
let database_url = env::var("TEST_DATABASE_URL")
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:5432/crank".to_owned());
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.connect(&database_url)
.await
.unwrap();
let schema = format!("test_mcp_transport_{}", Uuid::now_v7().simple());
admin_pool
.execute(sqlx::query(&format!("create schema {schema}")))
.await
.unwrap();
let connect_options = format!("{database_url}?options=-csearch_path%3D{schema}")
.parse::<PgConnectOptions>()
.unwrap();
let store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options.clone(),
PostgresPoolConfig::default(),
)
.await
.unwrap();
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
timestamp("2026-05-01T10:00:00Z"),
Some(timestamp("2026-05-01T10:00:01Z")),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
let pool = PgPoolOptions::new()
.max_connections(1)
.connect_with(connect_options)
.await
.unwrap();
let remaining = sqlx::query_scalar::<_, i64>(
"select count(*) from mcp_transport_sessions where id = $1",
)
.bind(&session_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining, 0);
}
}