crank/crates/crank-community-mcp/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)
}