rmcp_server_kit/
metrics.rs

1//! Prometheus metrics for MCP servers.
2//!
3//! Provides a shared [`crate::metrics::McpMetrics`] registry with standard HTTP counters.
4//! The transport layer exposes these via a `/metrics` endpoint on a
5//! dedicated listener when `metrics_enabled` is true.
6//!
7//! # Public surface and the `prometheus` crate
8//!
9//! [`crate::metrics::McpMetrics::registry`] and the `IntCounterVec` / `HistogramVec` fields are
10//! intentionally exposed so downstream crates can register additional custom
11//! collectors against the same registry. This re-exports the [`prometheus`]
12//! crate types as part of `rmcp-server-kit`'s public API; pin the same major version to
13//! avoid type-identity mismatches when registering custom metrics.
14
15use std::sync::Arc;
16
17use prometheus::{
18    Encoder, HistogramOpts, HistogramVec, IntCounterVec, Registry, TextEncoder, opts,
19};
20
21use crate::error::RmcpServerKitError;
22
23/// Default Prometheus histogram buckets for HTTP request latency
24/// (seconds). Tuned for low-latency service work: sub-millisecond
25/// through five seconds, covering health-check fast paths up to slow
26/// outbound dependencies. Operators that need different buckets can
27/// register their own histogram against
28/// [`McpMetrics::registry`].
29const HTTP_DURATION_BUCKETS: &[f64] = &[
30    0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0,
31];
32
33/// Collected Prometheus metrics for an MCP server.
34#[derive(Clone, Debug)]
35#[non_exhaustive]
36pub struct McpMetrics {
37    /// Prometheus registry holding all counters and histograms.
38    pub registry: Registry,
39    /// Total HTTP requests by method, path, and status code.
40    pub http_requests_total: IntCounterVec,
41    /// HTTP request duration in seconds by method and path.
42    pub http_request_duration_seconds: HistogramVec,
43    /// Rate-limiter denials by limiter. Label `limiter` is one of
44    /// `tool`, `auth_pre`, `auth_post`, `extra_route` - matching the
45    /// four built-in per-IP limiters. Incremented at each deny site
46    /// alongside the existing warn-level log.
47    pub rate_limited_total: IntCounterVec,
48}
49
50impl McpMetrics {
51    /// Create a new metrics registry with default MCP counters.
52    ///
53    /// # Errors
54    ///
55    /// Returns [`RmcpServerKitError::Metrics`] if counter registration fails (should
56    /// not happen unless duplicate registrations occur).
57    pub fn new() -> Result<Self, RmcpServerKitError> {
58        let registry = Registry::new();
59
60        let http_requests_total = IntCounterVec::new(
61            opts!("rmcp_server_kit_http_requests_total", "Total HTTP requests"),
62            &["method", "path", "status"],
63        )
64        .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
65        registry
66            .register(Box::new(http_requests_total.clone()))
67            .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
68
69        let http_request_duration_seconds = HistogramVec::new(
70            HistogramOpts::new(
71                "rmcp_server_kit_http_request_duration_seconds",
72                "HTTP request duration in seconds",
73            )
74            .buckets(HTTP_DURATION_BUCKETS.to_vec()),
75            &["method", "path"],
76        )
77        .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
78        registry
79            .register(Box::new(http_request_duration_seconds.clone()))
80            .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
81
82        let rate_limited_total = IntCounterVec::new(
83            opts!(
84                "rmcp_server_kit_rate_limited_total",
85                "Rate-limiter denials by limiter"
86            ),
87            &["limiter"],
88        )
89        .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
90        registry
91            .register(Box::new(rate_limited_total.clone()))
92            .map_err(|e| RmcpServerKitError::Metrics(e.to_string()))?;
93
94        Ok(Self {
95            registry,
96            http_requests_total,
97            http_request_duration_seconds,
98            rate_limited_total,
99        })
100    }
101
102    /// Encode all collected metrics as Prometheus text format.
103    #[must_use]
104    pub fn encode(&self) -> String {
105        let encoder = TextEncoder::new();
106        let metric_families = self.registry.gather();
107        let mut buf = Vec::new();
108        if let Err(e) = encoder.encode(&metric_families, &mut buf) {
109            tracing::warn!(error = %e, "prometheus encode failed");
110            return String::new();
111        }
112        // TextEncoder always produces valid UTF-8; fall back to empty on
113        // the near-impossible chance it doesn't.
114        String::from_utf8(buf).unwrap_or_default()
115    }
116}
117
118/// Increment the rate-limiter deny counter for `limiter`, if the shared
119/// [`McpMetrics`] handle is present in the request extensions.
120///
121/// The handle is inserted by the transport's metrics middleware (the
122/// outermost layer on the merged router) only when `metrics_enabled` is
123/// true; absent the extension this is a no-op, so deny sites behave
124/// identically with metrics disabled. `limiter` is one of `tool`,
125/// `auth_pre`, `auth_post`, `extra_route`.
126pub(crate) fn record_rate_limit_deny(ext: &axum::http::Extensions, limiter: &str) {
127    if let Some(m) = ext.get::<Arc<McpMetrics>>() {
128        m.rate_limited_total.with_label_values(&[limiter]).inc();
129    }
130}
131
132/// Spawn a dedicated HTTP listener that serves Prometheus metrics on `/metrics`.
133///
134/// The listener exits and releases the bound port when `shutdown` is
135/// cancelled, keeping the metrics endpoint tied to the parent server's
136/// graceful-shutdown lifecycle (M7).
137///
138/// # Errors
139///
140/// Returns [`RmcpServerKitError::Startup`] if the TCP listener cannot bind or the
141/// underlying axum server fails.
142// cancel-safe: the parent server cancels via `shutdown.cancelled()` inside
143// axum graceful shutdown; dropping this future directly only drops the
144// listener/app, with no metrics registry mutation or detached work.
145pub async fn serve_metrics(
146    bind: String,
147    metrics: Arc<McpMetrics>,
148    shutdown: tokio_util::sync::CancellationToken,
149) -> Result<(), RmcpServerKitError> {
150    let app = axum::Router::new().route(
151        "/metrics",
152        axum::routing::get(move || {
153            let m = Arc::clone(&metrics);
154            async move { m.encode() }
155        }),
156    );
157
158    let listener = tokio::net::TcpListener::bind(&bind)
159        .await
160        .map_err(|e| RmcpServerKitError::Startup(format!("metrics bind {bind}: {e}")))?;
161    tracing::info!("metrics endpoint listening on http://{bind}/metrics");
162    axum::serve(listener, app)
163        .with_graceful_shutdown(async move { shutdown.cancelled().await })
164        .await
165        .map_err(|e| RmcpServerKitError::Startup(format!("metrics serve: {e}")))?;
166    Ok(())
167}
168
169#[cfg(test)]
170mod tests {
171    #![allow(
172        clippy::unwrap_used,
173        clippy::expect_used,
174        clippy::panic,
175        clippy::indexing_slicing,
176        clippy::unwrap_in_result,
177        clippy::print_stdout,
178        clippy::print_stderr,
179        reason = "test-only relaxations; production code uses ? and tracing"
180    )]
181    use super::*;
182
183    #[test]
184    fn new_creates_registry_with_counters() {
185        let m = McpMetrics::new().unwrap();
186        // Incrementing a counter should make it appear in gather output.
187        m.http_requests_total
188            .with_label_values(&["GET", "/test", "200"])
189            .inc();
190        m.http_request_duration_seconds
191            .with_label_values(&["GET", "/test"])
192            .observe(0.1);
193        assert_eq!(m.registry.gather().len(), 2);
194    }
195
196    #[test]
197    fn encode_empty_registry() {
198        let m = McpMetrics::new().unwrap();
199        let output = m.encode();
200        // Empty counters/histograms produce no samples but the output is valid.
201        assert!(output.is_empty() || output.contains("rmcp_server_kit_"));
202    }
203
204    #[test]
205    fn counter_increment_shows_in_encode() {
206        let m = McpMetrics::new().unwrap();
207        m.http_requests_total
208            .with_label_values(&["GET", "/healthz", "200"])
209            .inc();
210        let output = m.encode();
211        assert!(output.contains("rmcp_server_kit_http_requests_total"));
212        assert!(output.contains("method=\"GET\""));
213        assert!(output.contains("path=\"/healthz\""));
214        assert!(output.contains("status=\"200\""));
215        assert!(output.contains(" 1")); // count = 1
216    }
217
218    #[test]
219    fn histogram_observe_shows_in_encode() {
220        let m = McpMetrics::new().unwrap();
221        m.http_request_duration_seconds
222            .with_label_values(&["POST", "/mcp"])
223            .observe(0.042);
224        let output = m.encode();
225        assert!(output.contains("rmcp_server_kit_http_request_duration_seconds"));
226        assert!(output.contains("method=\"POST\""));
227        assert!(output.contains("path=\"/mcp\""));
228    }
229
230    #[test]
231    fn multiple_increments_accumulate() {
232        let m = McpMetrics::new().unwrap();
233        let counter = m
234            .http_requests_total
235            .with_label_values(&["POST", "/mcp", "200"]);
236        counter.inc();
237        counter.inc();
238        counter.inc();
239        let output = m.encode();
240        assert!(output.contains(" 3")); // count = 3
241    }
242
243    #[test]
244    fn clone_shares_registry() {
245        let m = McpMetrics::new().unwrap();
246        let m2 = m.clone();
247        m.http_requests_total
248            .with_label_values(&["GET", "/test", "200"])
249            .inc();
250        // The clone should see the same counter value.
251        let output = m2.encode();
252        assert!(output.contains(" 1"));
253    }
254
255    #[test]
256    fn rate_limited_counter_registers_and_encodes() {
257        let m = McpMetrics::new().unwrap();
258        m.rate_limited_total.with_label_values(&["tool"]).inc();
259        let output = m.encode();
260        assert!(output.contains("rmcp_server_kit_rate_limited_total"));
261        assert!(output.contains("limiter=\"tool\""));
262        assert!(output.contains(" 1"));
263    }
264
265    #[test]
266    fn record_rate_limit_deny_increments_via_extension() {
267        let m = Arc::new(McpMetrics::new().unwrap());
268        let mut ext = axum::http::Extensions::new();
269        ext.insert(Arc::clone(&m));
270        record_rate_limit_deny(&ext, "auth_pre");
271        record_rate_limit_deny(&ext, "auth_pre");
272        assert_eq!(
273            m.rate_limited_total.with_label_values(&["auth_pre"]).get(),
274            2
275        );
276        // Absent handle: silent no-op (metrics disabled path).
277        let empty = axum::http::Extensions::new();
278        record_rate_limit_deny(&empty, "auth_pre");
279        assert_eq!(
280            m.rate_limited_total.with_label_values(&["auth_pre"]).get(),
281            2
282        );
283    }
284
285    // M7 regression: cancelling the shutdown token must release the
286    // metrics listener's bound port so a subsequent bind to the same
287    // address succeeds. Prior to M7 the metrics endpoint ran without
288    // graceful_shutdown wiring and would leak the port until process
289    // exit.
290    #[tokio::test]
291    async fn serve_metrics_releases_port_on_shutdown() {
292        // Pick an ephemeral port, then drop the probe so serve_metrics
293        // can claim it.
294        let probe = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
295        let addr = probe.local_addr().unwrap();
296        drop(probe);
297
298        let metrics = Arc::new(McpMetrics::new().unwrap());
299        let shutdown = tokio_util::sync::CancellationToken::new();
300        let handle = tokio::spawn(serve_metrics(
301            addr.to_string(),
302            Arc::clone(&metrics),
303            shutdown.clone(),
304        ));
305
306        // Wait until the listener is actually accepting connections.
307        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
308        loop {
309            if tokio::net::TcpStream::connect(addr).await.is_ok() {
310                break;
311            }
312            assert!(
313                std::time::Instant::now() < deadline,
314                "metrics listener never accepted on {addr}"
315            );
316            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
317        }
318
319        // Cancel and await graceful shutdown.
320        shutdown.cancel();
321        let join = tokio::time::timeout(std::time::Duration::from_secs(5), handle)
322            .await
323            .expect("serve_metrics did not return within timeout");
324        join.expect("join error")
325            .expect("serve_metrics returned Err");
326
327        // Port must be immediately rebindable.
328        let rebind = tokio::net::TcpListener::bind(addr)
329            .await
330            .expect("port not released after shutdown");
331        drop(rebind);
332    }
333}
rmcp_server_kit/metrics.rs

rmcp_server_kit/
metrics.rs
