diff --git a/TASKS.md b/TASKS.md index f7c3dca..0900790 100644 --- a/TASKS.md +++ b/TASKS.md @@ -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 diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index b3bd53d..2b9c2b2 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -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 { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index dcb240e..b90a3b1 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -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, 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, + 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(®istry, "sales-expired-session", vec![]).await; + let api_key = create_platform_api_key( + ®istry, + "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; diff --git a/apps/mcp-server/src/session.rs b/apps/mcp-server/src/session.rs index 6e3b5e1..06b1ebf 100644 --- a/apps/mcp-server/src/session.rs +++ b/apps/mcp-server/src/session.rs @@ -38,6 +38,7 @@ pub trait TransportSessionStore: Send + Sync { workspace_slug: &str, agent_slug: &str, now: OffsetDateTime, + expires_at: Option, ) -> Result; async fn get(&self, session_id: &str) -> Result, SessionStoreError>; @@ -100,6 +101,7 @@ impl TransportSessionStore for InMemorySessionStore { workspace_slug: &str, agent_slug: &str, now: OffsetDateTime, + expires_at: Option, ) -> Result { 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, 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, ) -> Result { 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::() + .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); + } }