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

mcp: add postgres transport session store

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 23:10:38 +00:00
parent 6c8be5c5d8
commit 77f83a1dc1
4 changed files with 331 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.sqlx/query-bdae00d03a55ba3b6d121578b081e970ff08427f9a3dd396b649996a9131260c.json
Generated
+64
View File
@@ -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"
}
apps/mcp-server/src/app.rs
+1 -20
View File
@@ -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<String>,
secret_crypto: SecretCrypto,
runtime: RuntimeExecutor,
api_rate_limiter: RequestRateLimiter,
sessions: SharedSessionStore,
) -> Router {
let state = Arc::new(AppState {
apps/mcp-server/src/main.rs
+11 -3
View File
@@ -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<dyn std::error::Error>> {
@@ -35,8 +35,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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()),
)
}
apps/mcp-server/src/session.rs
+255 -10
View File
@@ -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<dyn TransportSessionStore>;
#[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<Self, SessionStoreError> {
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<RwLock<HashMap<String, SessionState>>>,
}
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<String, SessionStoreError> {
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<Option<SessionState>, 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<bool, SessionStoreError> {
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<bool, SessionStoreError> {
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::<PgConnectOptions>()
.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");
}
}