crank/crates/crank-community-mcp/src/app/metrics.rs

use std::sync::Arc;

use axum::http::StatusCode;
use serde_json::Value;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};

use crate::{
    jsonrpc::{is_notification, is_response, method_name},
    transport::ResponseMode,
};

pub(super) struct McpRequestMetrics {
    method: &'static str,
    response_mode: &'static str,
    outcome: &'static str,
}

impl McpRequestMetrics {
    pub(super) fn new(message: &Value) -> Self {
        Self {
            method: normalized_mcp_method(message),
            response_mode: "unknown",
            outcome: "rejected",
        }
    }

    pub(super) fn set_response_mode(&mut self, mode: ResponseMode) -> ResponseMode {
        self.response_mode = match mode {
            ResponseMode::Json => "json",
            ResponseMode::Sse => "sse",
        };
        mode
    }

    pub(super) fn complete(&mut self, status: StatusCode) {
        self.outcome = match status.as_u16() {
            200..=299 => "success",
            400..=499 => "client_error",
            500..=599 => "server_error",
            _ => "other",
        };
    }
}

impl Drop for McpRequestMetrics {
    fn drop(&mut self) {
        ::metrics::counter!(
            "crank_mcp_requests_total",
            "method" => self.method,
            "response_mode" => self.response_mode,
            "outcome" => self.outcome
        )
        .increment(1);
    }
}

pub(super) fn normalized_mcp_method(message: &Value) -> &'static str {
    match method_name(message) {
        Some("initialize") => "initialize",
        Some("notifications/initialized") => "initialized",
        Some("ping") => "ping",
        Some("tools/list") => "tools_list",
        Some("tools/call") => "tools_call",
        Some(_) if is_notification(message) => "notification",
        Some(_) => "unsupported",
        None if is_response(message) => "response",
        None => "invalid",
    }
}

pub(super) struct ActiveSessionGuard {
    _permit: OwnedSemaphorePermit,
}

impl ActiveSessionGuard {
    pub(super) fn try_acquire(slots: &Arc<Semaphore>) -> Result<Self, ()> {
        let permit = Arc::clone(slots).try_acquire_owned().map_err(|_| {
            ::metrics::counter!(
                "crank_runtime_limit_rejections_total",
                "stage" => "mcp_session"
            )
            .increment(1);
        })?;
        ::metrics::gauge!("crank_mcp_active_sessions").increment(1.0);
        Ok(Self { _permit: permit })
    }
}

impl Drop for ActiveSessionGuard {
    fn drop(&mut self) {
        ::metrics::gauge!("crank_mcp_active_sessions").decrement(1.0);
    }
}