crank/apps/mcp-server/src/session.rs

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;
use uuid::Uuid;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionState {
    pub id: String,
    pub protocol_version: String,
    pub initialized: bool,
    pub workspace_slug: String,
    pub agent_slug: String,
    pub created_at: OffsetDateTime,
    pub updated_at: OffsetDateTime,
    pub expires_at: Option<OffsetDateTime>,
}

#[derive(Debug, Error)]
#[error("transport session store is unavailable: {details}")]
pub struct SessionStoreError {
    pub details: String,
}

#[async_trait]
pub trait TransportSessionStore: Send + Sync {
    async fn create(
        &self,
        protocol_version: &str,
        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>;

    async fn mark_initialized(
        &self,
        session_id: &str,
        now: OffsetDateTime,
    ) -> Result<bool, SessionStoreError>;

    async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
}

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>>>,
}

#[async_trait]
impl TransportSessionStore for InMemorySessionStore {
    async fn create(
        &self,
        protocol_version: &str,
        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;

        guard.insert(
            session_id.clone(),
            SessionState {
                id: session_id.clone(),
                protocol_version: protocol_version.to_owned(),
                initialized: false,
                workspace_slug: workspace_slug.to_owned(),
                agent_slug: agent_slug.to_owned(),
                created_at: now,
                updated_at: now,
                expires_at,
            },
        );

        Ok(session_id)
    }

    async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
        {
            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(
        &self,
        session_id: &str,
        now: OffsetDateTime,
    ) -> Result<bool, SessionStoreError> {
        let mut guard = self.inner.write().await;

        if let Some(session) = guard.get_mut(session_id) {
            session.initialized = true;
            session.updated_at = now;
            return Ok(true);
        }

        Ok(false)
    }

    async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
        let mut guard = self.inner.write().await;
        Ok(guard.remove(session_id).is_some())
    }
}

#[async_trait]
impl TransportSessionStore for PostgresTransportSessionStore {
    async fn create(
        &self,
        protocol_version: &str,
        workspace_slug: &str,
        agent_slug: &str,
        now: OffsetDateTime,
        expires_at: Option<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, $6::timestamptz
             )",
        )
        .bind(&session_id)
        .bind(protocol_version)
        .bind(workspace_slug)
        .bind(agent_slug)
        .bind(now)
        .bind(expires_at)
        .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(),
        })?;

        let Some(session) = 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,
        }) 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(
        &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(())
}

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;

    use sqlx::{
        Executor,
        postgres::{PgConnectOptions, PgPoolOptions},
    };
    use time::{OffsetDateTime, 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()
    }

    fn truncate_to_micros(value: OffsetDateTime) -> OffsetDateTime {
        value
            .replace_nanosecond((value.nanosecond() / 1_000) * 1_000)
            .unwrap()
    }

    #[tokio::test]
    async fn creates_and_reads_transport_sessions() {
        let store = InMemorySessionStore::default();
        let created_at = timestamp("2026-05-01T10:00:00Z");

        let session_id = store
            .create("2025-11-25", "default", "sales", created_at, None)
            .await
            .unwrap();
        let session = store.get(&session_id).await.unwrap().unwrap();

        assert_eq!(session.id, session_id);
        assert_eq!(session.protocol_version, "2025-11-25");
        assert_eq!(session.workspace_slug, "default");
        assert_eq!(session.agent_slug, "sales");
        assert!(!session.initialized);
        assert_eq!(session.created_at, created_at);
        assert_eq!(session.updated_at, created_at);
    }

    #[tokio::test]
    async fn marks_transport_sessions_initialized() {
        let store = InMemorySessionStore::default();
        let created_at = timestamp("2026-05-01T10:00:00Z");
        let initialized_at = timestamp("2026-05-01T10:00:05Z");

        let session_id = store
            .create("2025-11-25", "default", "sales", created_at, None)
            .await
            .unwrap();

        assert!(
            store
                .mark_initialized(&session_id, initialized_at)
                .await
                .unwrap()
        );

        let session = store.get(&session_id).await.unwrap().unwrap();
        assert!(session.initialized);
        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 {
            details: "postgres is unavailable".to_owned(),
        };

        assert_eq!(
            error.to_string(),
            "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 = OffsetDateTime::now_utc() - time::Duration::hours(1);
        let initialized_at = created_at + time::Duration::seconds(5);
        let session_id = store_a
            .create(
                "2025-11-25",
                "default",
                "sales",
                created_at,
                Some(created_at + time::Duration::days(30)),
            )
            .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, truncate_to_micros(initialized_at));
        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);
    }
}