From 77f83a1dc16827650c730c198f4409bbf7c3b099 Mon Sep 17 00:00:00 2001 From: "a.tolmachev" Date: Fri, 1 May 2026 23:10:38 +0000 Subject: [PATCH] mcp: add postgres transport session store --- ...1e970ff08427f9a3dd396b649996a9131260c.json | 64 +++++ apps/mcp-server/src/app.rs | 21 +- apps/mcp-server/src/main.rs | 14 +- apps/mcp-server/src/session.rs | 265 +++++++++++++++++- 4 files changed, 331 insertions(+), 33 deletions(-) create mode 100644 .sqlx/query-bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c.json diff --git a/.sqlx/query-bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c.json b/.sqlx/query-bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c.json new file mode 100644 index 0000000..8762d72 --- /dev/null +++ b/.sqlx/query-bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c.json @@ -0,0 +1,64 @@ +{ + "db_name": "PostgreSQL", + "query": "select\n id,\n protocol_version,\n initialized,\n workspace_slug,\n agent_slug,\n created_at as \"created_at!: OffsetDateTime\",\n updated_at as \"updated_at!: OffsetDateTime\",\n expires_at as \"expires_at: OffsetDateTime\"\n from mcp_transport_sessions\n where id = $1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": "Text" + }, + { + "ordinal": 1, + "name": "protocol_version", + "type_info": "Text" + }, + { + "ordinal": 2, + "name": "initialized", + "type_info": "Bool" + }, + { + "ordinal": 3, + "name": "workspace_slug", + "type_info": "Text" + }, + { + "ordinal": 4, + "name": "agent_slug", + "type_info": "Text" + }, + { + "ordinal": 5, + "name": "created_at!: OffsetDateTime", + "type_info": "Timestamptz" + }, + { + "ordinal": 6, + "name": "updated_at!: OffsetDateTime", + "type_info": "Timestamptz" + }, + { + "ordinal": 7, + "name": "expires_at: OffsetDateTime", + "type_info": "Timestamptz" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + true + ] + }, + "hash": "bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c" +} diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 7c7d0b6..b3bd53d 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -47,7 +47,7 @@ use crate::{ is_response, jsonrpc_error, jsonrpc_result, method_name, negotiated_protocol_version, params, request_id, }, - session::{InMemorySessionStore, SessionState, SharedSessionStore}, + session::{SessionState, SharedSessionStore}, }; const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id"; @@ -131,25 +131,6 @@ 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, - secret_crypto: SecretCrypto, - runtime: RuntimeExecutor, - api_rate_limiter: RequestRateLimiter, sessions: SharedSessionStore, ) -> Router { let state = Arc::new(AppState { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 90f3af0..dcb240e 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -13,7 +13,7 @@ use sqlx::postgres::PgConnectOptions; use tokio::net::TcpListener; use tracing::info; -use crate::app::build_app; +use crate::{app::build_app, session::PostgresTransportSessionStore}; #[tokio::main] async fn main() -> Result<(), Box> { @@ -35,8 +35,14 @@ async fn main() -> Result<(), Box> { let pool_config = PostgresPoolConfig::from_env()?; let runtime_limits = RuntimeLimits::from_env()?; let api_rate_limit = mcp_api_rate_limit_config_from_env()?; + let database_options = database_options_from_env()?; let registry = PostgresRegistry::connect_with_options_and_pool_config( - database_options_from_env()?, + database_options.clone(), + pool_config, + ) + .await?; + let session_store = PostgresTransportSessionStore::connect_with_options_and_pool_config( + database_options, pool_config, ) .await?; @@ -49,6 +55,7 @@ async fn main() -> Result<(), Box> { secret_crypto, runtime, RequestRateLimiter::new(api_rate_limit), + std::sync::Arc::new(session_store), ); let listener = TcpListener::bind(socket_addr).await?; @@ -151,7 +158,7 @@ mod tests { use tokio::time::sleep; use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*}; - use crate::app::build_app; + use crate::{app::build_app, session::InMemorySessionStore}; fn test_workspace_id() -> WorkspaceId { WorkspaceId::new("ws_default") @@ -223,6 +230,7 @@ mod tests { SecretCrypto::new("test-master-key").unwrap(), RuntimeExecutor::new(), RequestRateLimiter::new(rate_limit_config), + std::sync::Arc::new(InMemorySessionStore::default()), ) } diff --git a/apps/mcp-server/src/session.rs b/apps/mcp-server/src/session.rs index bc5f35d..6e3b5e1 100644 --- a/apps/mcp-server/src/session.rs +++ b/apps/mcp-server/src/session.rs @@ -1,6 +1,12 @@ use std::{collections::HashMap, sync::Arc}; use async_trait::async_trait; +use crank_registry::PostgresPoolConfig; +use sqlx::{ + PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, + query, +}; use thiserror::Error; use time::OffsetDateTime; use tokio::sync::RwLock; @@ -47,17 +53,45 @@ pub trait TransportSessionStore: Send + Sync { pub type SharedSessionStore = Arc; +#[derive(Clone, Debug)] +pub struct PostgresTransportSessionStore { + pool: PgPool, +} + +impl PostgresTransportSessionStore { + pub async fn connect_with_options_and_pool_config( + connect_options: PgConnectOptions, + pool_config: PostgresPoolConfig, + ) -> Result { + let pool = PgPoolOptions::new() + .max_connections(pool_config.max_connections) + .min_connections(pool_config.min_connections) + .acquire_timeout(std::time::Duration::from_millis( + pool_config.acquire_timeout_ms, + )) + .idle_timeout(Some(std::time::Duration::from_millis( + pool_config.idle_timeout_ms, + ))) + .max_lifetime(Some(std::time::Duration::from_millis( + pool_config.max_lifetime_ms, + ))) + .connect_with(connect_options) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + apply_postgres_migrations(&pool).await?; + + Ok(Self { pool }) + } +} + #[derive(Clone, Default)] pub struct InMemorySessionStore { inner: Arc>>, } -impl InMemorySessionStore { - pub fn new() -> Self { - Self::default() - } -} - #[async_trait] impl TransportSessionStore for InMemorySessionStore { async fn create( @@ -114,11 +148,169 @@ impl TransportSessionStore for InMemorySessionStore { } } +#[async_trait] +impl TransportSessionStore for PostgresTransportSessionStore { + async fn create( + &self, + protocol_version: &str, + workspace_slug: &str, + agent_slug: &str, + now: OffsetDateTime, + ) -> Result { + let session_id = Uuid::now_v7().to_string(); + query( + "insert into mcp_transport_sessions ( + id, + protocol_version, + initialized, + workspace_slug, + agent_slug, + created_at, + updated_at, + expires_at + ) values ( + $1, $2, false, $3, $4, $5::timestamptz, $5::timestamptz, null + )", + ) + .bind(&session_id) + .bind(protocol_version) + .bind(workspace_slug) + .bind(agent_slug) + .bind(now) + .execute(&self.pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + Ok(session_id) + } + + async fn get(&self, session_id: &str) -> Result, SessionStoreError> { + let row = sqlx::query!( + "select + id, + protocol_version, + initialized, + workspace_slug, + agent_slug, + created_at as \"created_at!: OffsetDateTime\", + updated_at as \"updated_at!: OffsetDateTime\", + expires_at as \"expires_at: OffsetDateTime\" + from mcp_transport_sessions + where id = $1", + session_id, + ) + .fetch_optional(&self.pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + Ok(row.map(|row| SessionState { + id: row.id, + protocol_version: row.protocol_version, + initialized: row.initialized, + workspace_slug: row.workspace_slug, + agent_slug: row.agent_slug, + created_at: row.created_at, + updated_at: row.updated_at, + expires_at: row.expires_at, + })) + } + + async fn mark_initialized( + &self, + session_id: &str, + now: OffsetDateTime, + ) -> Result { + let result = query( + "update mcp_transport_sessions + set initialized = true, + updated_at = $2::timestamptz + where id = $1", + ) + .bind(session_id) + .bind(now) + .execute(&self.pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + Ok(result.rows_affected() > 0) + } + + async fn delete(&self, session_id: &str) -> Result { + let result = query("delete from mcp_transport_sessions where id = $1") + .bind(session_id) + .execute(&self.pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + Ok(result.rows_affected() > 0) + } +} + +async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreError> { + query( + "create table if not exists mcp_transport_sessions ( + id text primary key, + protocol_version text not null, + initialized boolean not null default false, + workspace_slug text not null, + agent_slug text not null, + created_at timestamptz not null, + updated_at timestamptz not null, + expires_at timestamptz null + )", + ) + .execute(pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + query( + "alter table mcp_transport_sessions add column if not exists expires_at timestamptz null", + ) + .execute(pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + query( + "create index if not exists mcp_transport_sessions_workspace_agent_idx + on mcp_transport_sessions(workspace_slug, agent_slug, updated_at desc)", + ) + .execute(pool) + .await + .map_err(|error| SessionStoreError { + details: error.to_string(), + })?; + + Ok(()) +} + #[cfg(test)] mod tests { - use time::format_description::well_known::Rfc3339; + use std::env; - use super::{InMemorySessionStore, SessionStoreError, TransportSessionStore}; + use sqlx::{ + Executor, + postgres::{PgConnectOptions, PgPoolOptions}, + }; + use time::format_description::well_known::Rfc3339; + use uuid::Uuid; + + use super::{ + InMemorySessionStore, PostgresTransportSessionStore, SessionStoreError, + TransportSessionStore, + }; + use crank_registry::PostgresPoolConfig; fn timestamp(value: &str) -> time::OffsetDateTime { time::OffsetDateTime::parse(value, &Rfc3339).unwrap() @@ -126,7 +318,7 @@ mod tests { #[tokio::test] async fn creates_and_reads_transport_sessions() { - let store = InMemorySessionStore::new(); + let store = InMemorySessionStore::default(); let created_at = timestamp("2026-05-01T10:00:00Z"); let session_id = store @@ -146,7 +338,7 @@ mod tests { #[tokio::test] async fn marks_transport_sessions_initialized() { - let store = InMemorySessionStore::new(); + let store = InMemorySessionStore::default(); let created_at = timestamp("2026-05-01T10:00:00Z"); let initialized_at = timestamp("2026-05-01T10:00:05Z"); @@ -178,4 +370,57 @@ mod tests { "transport session store is unavailable: postgres is unavailable" ); } + + #[tokio::test] + async fn postgres_transport_sessions_survive_store_reconnect() { + 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 pool_config = PostgresPoolConfig::default(); + let store_a = PostgresTransportSessionStore::connect_with_options_and_pool_config( + connect_options.clone(), + pool_config, + ) + .await + .unwrap(); + + 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) + .await + .unwrap(); + store_a + .mark_initialized(&session_id, initialized_at) + .await + .unwrap(); + + let store_b = PostgresTransportSessionStore::connect_with_options_and_pool_config( + connect_options, + pool_config, + ) + .await + .unwrap(); + let session = store_b.get(&session_id).await.unwrap().unwrap(); + + assert_eq!(session.id, session_id); + assert!(session.initialized); + assert_eq!(session.updated_at, initialized_at); + assert_eq!(session.workspace_slug, "default"); + assert_eq!(session.agent_slug, "sales"); + } }