crank/crates/crank-community-mcp/tests/session_metrics.rs

use std::{sync::Arc, time::Duration};

use crank_community_mcp::session::{
    ActiveSessionMetrics, InMemorySessionStore, SharedSessionStore, TransportSessionStore,
    spawn_session_cleanup,
};
use metrics_util::debugging::{DebugValue, DebuggingRecorder, Snapshotter};
use time::OffsetDateTime;

#[tokio::test]
async fn active_session_sampler_tracks_create_delete_and_expiry() {
    let recorder = DebuggingRecorder::new();
    let snapshotter = recorder.snapshotter();
    recorder
        .install()
        .expect("isolated integration test recorder");

    let store = Arc::new(InMemorySessionStore::default());
    let sessions: SharedSessionStore = store.clone();
    let sampler = ActiveSessionMetrics::start(Arc::clone(&sessions));
    assert_gauge(&sampler, &snapshotter, 0.0).await;

    let now = OffsetDateTime::now_utc();
    store
        .create(
            "2025-11-25",
            "default",
            "agent",
            false,
            now,
            Some(now + time::Duration::milliseconds(100)),
        )
        .await
        .unwrap();
    assert_gauge(&sampler, &snapshotter, 1.0).await;

    spawn_session_cleanup(sessions, sampler, Duration::from_millis(10));
    assert_gauge_without_refresh(&snapshotter, 0.0).await;
}

async fn assert_gauge_without_refresh(snapshotter: &Snapshotter, expected: f64) {
    for _ in 0..100 {
        tokio::time::sleep(Duration::from_millis(10)).await;
        if gauge_value(snapshotter) == Some(expected) {
            return;
        }
    }
    panic!("active session gauge did not become {expected} through cleanup");
}

async fn assert_gauge(sampler: &ActiveSessionMetrics, snapshotter: &Snapshotter, expected: f64) {
    for _ in 0..50 {
        sampler.refresh();
        tokio::time::sleep(Duration::from_millis(10)).await;
        if gauge_value(snapshotter) == Some(expected) {
            return;
        }
    }
    panic!("active session gauge did not become {expected}");
}

fn gauge_value(snapshotter: &Snapshotter) -> Option<f64> {
    snapshotter
        .snapshot()
        .into_vec()
        .into_iter()
        .find_map(|(key, _, _, value)| {
            if key.key().name() != "crank_mcp_active_sessions" {
                return None;
            }
            match value {
                DebugValue::Gauge(value) => Some(value.into_inner()),
                _ => None,
            }
        })
}