crank/crates/crank-metrics/tests/contract.rs

use std::{collections::BTreeSet, time::Duration};

use crank_metrics::{
    CacheOutcome, ConfirmationOutcome, DbPoolState, Exporter, HttpMethod, HttpRoute,
    HttpStatusClass, IdempotencyOutcome, InFlightGuard, InvocationSource, LimitStage, McpMethod,
    McpOutcome, McpResponseMode, SignalType, ToolErrorKind, ToolInvocationMetrics, ToolOutcome,
    UpstreamOperationKind, UpstreamOutcome, UpstreamRequestMetrics, initialize_gauges,
    metric_schema, record_cache_outcome, record_confirmation_outcome, record_export_failure,
    record_http_request, record_idempotency_outcome, record_invocation_history_lost,
    record_limit_rejection, record_mcp_request, record_tool_invocation, record_upstream_request,
    set_catalog, set_db_pool_connections, set_mcp_active_sessions,
};
use metrics_util::debugging::DebuggingRecorder;

#[test]
fn schema_names_are_unique_and_labels_are_closed() {
    let schema = metric_schema();
    let names = schema
        .iter()
        .map(|metric| metric.name)
        .collect::<BTreeSet<_>>();

    assert_eq!(names.len(), schema.len());
    assert!(names.contains("crank_runtime_cache_total"));
    assert!(names.contains("crank_idempotency_total"));
    assert!(names.contains("crank_confirmation_total"));
    assert!(names.contains("crank_mcp_active_streams"));

    let allowed_labels = [
        "route",
        "method",
        "status_class",
        "response_mode",
        "outcome",
        "source",
        "error_kind",
        "operation_kind",
        "stage",
        "state",
        "signal_type",
        "exporter",
    ];
    for metric in schema {
        assert!(
            metric
                .labels
                .iter()
                .all(|label| allowed_labels.contains(label)),
            "{} contains an unapproved label",
            metric.name
        );
    }
}

#[test]
fn untrusted_http_values_collapse_to_closed_variants() {
    assert_eq!(
        HttpRoute::from_matched_path("/api/admin/workspaces/{workspace_id}/operations").as_str(),
        "/api/admin/workspaces/{workspace_id}/operations"
    );
    assert_eq!(
        HttpRoute::from_matched_path("/customer-controlled"),
        HttpRoute::unmatched()
    );
    assert_eq!(HttpMethod::classify("CUSTOM"), HttpMethod::Other);
    assert_eq!(HttpStatusClass::from_status(999), HttpStatusClass::Other);
}

#[test]
fn application_outcomes_are_distinct_from_transport_success() {
    assert_eq!(McpOutcome::from_http_status(200), McpOutcome::Success);
    assert_ne!(McpOutcome::JsonRpcError, McpOutcome::Success);
    assert_ne!(McpOutcome::ToolError, McpOutcome::Success);

    assert_eq!(CacheOutcome::Hit.as_str(), "hit");
    assert_eq!(CacheOutcome::Miss.as_str(), "miss");
    assert_eq!(IdempotencyOutcome::Replay.as_str(), "replay");
    assert_eq!(IdempotencyOutcome::Conflict.as_str(), "conflict");
    assert_eq!(ConfirmationOutcome::Approved.as_str(), "approved");
    assert_eq!(ConfirmationOutcome::Required.as_str(), "required");
}

#[test]
fn dropped_product_guards_record_aborted_outcomes() {
    let recorder = DebuggingRecorder::new();
    let snapshotter = recorder.snapshotter();

    metrics::with_local_recorder(&recorder, || {
        drop(ToolInvocationMetrics::start(
            InvocationSource::AgentToolCall,
        ));
        drop(UpstreamRequestMetrics::start(UpstreamOperationKind::Rest));
    });

    let snapshot = snapshotter.snapshot().into_vec();
    assert!(snapshot.iter().any(|(key, _, _, _)| {
        key.key().name() == "crank_tool_invocations_total"
            && has_label(key.key(), "outcome", "aborted")
            && has_label(key.key(), "error_kind", "aborted")
    }));
    assert!(snapshot.iter().any(|(key, _, _, _)| {
        key.key().name() == "crank_upstream_requests_total"
            && has_label(key.key(), "outcome", "aborted")
    }));
}

fn has_label(key: &metrics::Key, name: &str, value: &str) -> bool {
    key.labels()
        .any(|label| label.key() == name && label.value() == value)
}

#[test]
fn diverse_calls_cannot_create_unbounded_series() {
    let recorder = DebuggingRecorder::new();
    let snapshotter = recorder.snapshotter();

    metrics::with_local_recorder(&recorder, || {
        for status in 200..300 {
            let route = format!("/customer-controlled/{status}");
            let method = format!("CUSTOM-{status}");
            record_http_request(
                HttpRoute::from_matched_path(&route),
                HttpMethod::classify(&method),
                HttpStatusClass::from_status(status),
                Duration::from_millis(u64::from(status)),
            );
            record_mcp_request(
                McpMethod::ToolsCall,
                McpResponseMode::Json,
                McpOutcome::ToolError,
            );
            record_tool_invocation(
                InvocationSource::AgentToolCall,
                ToolOutcome::Error,
                ToolErrorKind::Mapping,
                Duration::from_millis(u64::from(status)),
            );
            record_cache_outcome(CacheOutcome::Miss);
            record_idempotency_outcome(IdempotencyOutcome::Replay);
            record_confirmation_outcome(ConfirmationOutcome::Required);
        }
    });

    let snapshot = snapshotter.snapshot().into_vec();
    let series = snapshot
        .iter()
        .map(|(key, _, _, _)| format!("{:?}", key.key()))
        .collect::<BTreeSet<_>>();

    assert_eq!(series.len(), 8);
    assert!(
        series
            .iter()
            .all(|series| !series.contains("customer-controlled"))
    );
}

#[test]
fn every_recording_api_matches_the_declared_schema() {
    let recorder = DebuggingRecorder::new();
    let snapshotter = recorder.snapshotter();

    metrics::with_local_recorder(&recorder, || {
        initialize_gauges();
        record_http_request(
            HttpRoute::from_matched_path("/api/auth/session"),
            HttpMethod::Get,
            HttpStatusClass::Success,
            Duration::from_millis(1),
        );
        record_mcp_request(
            McpMethod::ToolsCall,
            McpResponseMode::Json,
            McpOutcome::Success,
        );
        set_mcp_active_sessions(1);
        let _stream = InFlightGuard::mcp_stream();
        let _runtime = InFlightGuard::runtime();
        record_tool_invocation(
            InvocationSource::AgentToolCall,
            ToolOutcome::Error,
            ToolErrorKind::Mapping,
            Duration::from_millis(1),
        );
        record_upstream_request(
            UpstreamOperationKind::Rest,
            UpstreamOutcome::Timeout,
            Duration::from_millis(1),
        );
        record_limit_rejection(LimitStage::Concurrency);
        record_cache_outcome(CacheOutcome::Hit);
        record_idempotency_outcome(IdempotencyOutcome::Replay);
        record_confirmation_outcome(ConfirmationOutcome::Required);
        set_db_pool_connections(DbPoolState::Idle, 1);
        set_catalog(1, 2, 3);
        record_invocation_history_lost();
        record_export_failure(SignalType::Trace, Exporter::Otlp);
    });

    let snapshot = snapshotter.snapshot().into_vec();
    let actual_names = snapshot
        .iter()
        .map(|(key, _, _, _)| key.key().name().to_owned())
        .collect::<BTreeSet<_>>();
    let expected_names = metric_schema()
        .iter()
        .map(|definition| definition.name.to_owned())
        .collect::<BTreeSet<_>>();

    assert_eq!(actual_names, expected_names);
    for definition in metric_schema() {
        let actual_label_sets = snapshot
            .iter()
            .filter(|(key, _, _, _)| key.key().name() == definition.name)
            .map(|(key, _, _, _)| {
                key.key()
                    .labels()
                    .map(|label| label.key().to_owned())
                    .collect::<BTreeSet<_>>()
            })
            .collect::<BTreeSet<_>>();
        let expected_labels = definition
            .labels
            .iter()
            .map(|label| (*label).to_owned())
            .collect::<BTreeSet<_>>();

        assert_eq!(
            actual_label_sets,
            BTreeSet::from([expected_labels]),
            "recording API for {} diverges from the schema",
            definition.name
        );
    }
}