crank/crates/crank-community-mcp/tests/integration/session.rs

use crank_community_mcp::session::{PostgresTransportSessionStore, TransportSessionStore};
use crank_registry::PostgresPoolConfig;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};

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 postgres_transport_sessions_survive_store_reconnect() {
    let database_url = crank_test_support::postgres_schema_url("test_mcp_transport").await;
    let connect_options = database_url.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",
            false,
            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 = crank_test_support::postgres_schema_url("test_mcp_transport").await;
    let connect_options = database_url.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",
            false,
            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);
}

#[tokio::test]
async fn postgres_transport_session_cleanup_removes_abandoned_expired_rows() {
    let database_url = crank_test_support::postgres_schema_url("test_mcp_cleanup").await;
    let store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
        database_url.parse::<PgConnectOptions>().unwrap(),
        PostgresPoolConfig::default(),
    )
    .await
    .unwrap();
    let now = OffsetDateTime::now_utc();

    store
        .create(
            "2025-11-25",
            "default",
            "sales",
            false,
            now - time::Duration::hours(2),
            Some(now - time::Duration::hours(1)),
        )
        .await
        .unwrap();
    let active = store
        .create(
            "2025-11-25",
            "default",
            "sales",
            false,
            now,
            Some(now + time::Duration::hours(1)),
        )
        .await
        .unwrap();

    assert_eq!(store.cleanup_expired(now).await.unwrap(), 1);
    assert_eq!(store.active_count(now).await.unwrap(), 1);
    assert!(store.get(&active).await.unwrap().is_some());

    assert!(store.delete(&active).await.unwrap());
    assert_eq!(store.active_count(now).await.unwrap(), 0);
}