crank/crates/crank-observability/tests/metrics_signals.rs

use std::{
    net::{IpAddr, Ipv4Addr, SocketAddr},
    time::Duration,
};

use crank_metrics::{
    CacheOutcome, ConfirmationOutcome, DbPoolState, Exporter, HttpMethod, HttpRoute,
    HttpStatusClass, IdempotencyOutcome, InFlightGuard, InvocationSource, LimitStage, McpMethod,
    McpOutcome, McpResponseMode, SignalType, ToolErrorKind, ToolOutcome, UpstreamOperationKind,
    UpstreamOutcome, 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 crank_observability::{MetricsConfig, ObservabilityConfig, RedactionLimits, ServiceIdentity};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

#[tokio::test]
async fn typed_product_signals_are_exposed_by_a_real_scrape() {
    let identity =
        ServiceIdentity::try_new("metrics-signals", "0.3.1", "test").expect("valid identity");
    let lifecycle = crank_observability::init(ObservabilityConfig::new(
        identity,
        "off",
        RedactionLimits::default(),
    ))
    .expect("observability lifecycle");
    let config = MetricsConfig::new(
        true,
        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
        None,
    )
    .expect("loopback metrics config");
    let server = lifecycle
        .metrics_surface(config)
        .bind()
        .await
        .expect("metrics listener");
    let address = server.local_addr().expect("listener address");
    let task = tokio::spawn(server.serve());

    record_http_request(
        HttpRoute::from_matched_path("/api/auth/session"),
        HttpMethod::Get,
        HttpStatusClass::Success,
        Duration::from_millis(7),
    );
    record_mcp_request(
        McpMethod::ToolsCall,
        McpResponseMode::Json,
        McpOutcome::ToolError,
    );
    set_mcp_active_sessions(2);
    let _stream = InFlightGuard::mcp_stream();
    let _runtime = InFlightGuard::runtime();
    record_tool_invocation(
        InvocationSource::AgentToolCall,
        ToolOutcome::Error,
        ToolErrorKind::Mapping,
        Duration::from_millis(11),
    );
    record_upstream_request(
        UpstreamOperationKind::Rest,
        UpstreamOutcome::Timeout,
        Duration::from_millis(13),
    );
    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, 3);
    set_db_pool_connections(DbPoolState::Used, 1);
    set_catalog(5, 1200, 1);
    record_invocation_history_lost();
    record_export_failure(SignalType::Trace, Exporter::Otlp);

    let response = raw_get(address, "/metrics").await;
    task.abort();
    let _ = task.await;

    assert!(response.starts_with("HTTP/1.1 200 OK"), "{response}");
    for metric in [
        "crank_http_requests_total",
        "crank_mcp_requests_total",
        "crank_mcp_active_sessions",
        "crank_mcp_active_streams",
        "crank_tool_invocations_total",
        "crank_upstream_requests_total",
        "crank_runtime_inflight",
        "crank_runtime_limit_rejections_total",
        "crank_runtime_cache_total",
        "crank_idempotency_total",
        "crank_confirmation_total",
        "crank_db_pool_connections",
        "crank_catalog_tools",
        "crank_invocation_history_lost_total",
        "crank_telemetry_export_failures_total",
    ] {
        assert!(response.contains(metric), "missing metric {metric}");
    }
    for expected_label in [
        "outcome=\"tool_error\"",
        "outcome=\"timeout\"",
        "outcome=\"hit\"",
        "outcome=\"replay\"",
        "outcome=\"required\"",
        "error_kind=\"mapping\"",
    ] {
        assert!(
            response.contains(expected_label),
            "missing label {expected_label}"
        );
    }
    assert!(!response.contains("customer-secret"));
}

async fn raw_get(address: SocketAddr, path: &str) -> String {
    let mut stream = tokio::net::TcpStream::connect(address)
        .await
        .expect("metrics listener connection");
    stream
        .write_all(
            format!("GET {path} HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n")
                .as_bytes(),
        )
        .await
        .expect("metrics request");
    let mut response = Vec::new();
    stream
        .read_to_end(&mut response)
        .await
        .expect("metrics response");
    String::from_utf8(response).expect("utf-8 response")
}