mcp: add postgres transport session store

2026-05-01 23:10:38 +00:00
parent 6c8be5c5d8
commit 77f83a1dc1
4 changed files with 331 additions and 33 deletions
@@ -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");
+    }
 }