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