github-ops
100 lines
3.1 KiB
Rust
100 lines
3.1 KiB
Rust
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);
|
|
}
|