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    ///
104    /// On encoder failure the response is a **non-empty, stable marker body**
105    /// rather than an empty string, and the failure is logged at ERROR: with
106    /// caller-registered collectors admitted, one malformed collector must not
107    /// silently blank an entire scrape.
108    #[must_use]
109    pub fn encode(&self) -> String {
110        let encoder = TextEncoder::new();
111        let metric_families = self.registry.gather();
112        let mut buf = Vec::new();
113        if let Err(error) = encoder.encode(&metric_families, &mut buf) {
114            return encode_failure_body(&error);
115        }
116        // TextEncoder always produces valid UTF-8; fall back to empty on
117        // the near-impossible chance it doesn't.
118        String::from_utf8(buf).unwrap_or_default()
119    }
120}
121
122/// Marker body served when Prometheus encoding fails. Public within the crate
123/// so tests can assert its stability; scrapers and alert rules can key on it.
124pub(crate) const ENCODE_FAILURE_MARKER: &str =
125    "# rmcp-server-kit: prometheus encoding failed - see server logs\n";
126
127/// Build the failure body and log the underlying encoder error at ERROR.
128///
129/// The text encoder's only failure mode is an IO error from the sink (it
130/// performs no content validation), so this branch is unreachable with the
131/// in-memory `Vec` sink [`McpMetrics::encode`] uses - it exists so that a
132/// future writer-backed path, or a prometheus version that adds validation,
133/// cannot serve an empty scrape.
134fn encode_failure_body(error: &prometheus::Error) -> String {
135    tracing::error!(error = %error, "prometheus encoding failed; serving error marker body");
136    ENCODE_FAILURE_MARKER.to_owned()
137}
138
139/// Increment the rate-limiter deny counter for `limiter`, if the shared
140/// [`McpMetrics`] handle is present in the request extensions.
141///
142/// The handle is inserted by the transport's metrics middleware (the
143/// outermost layer on the merged router) only when `metrics_enabled` is
144/// true; absent the extension this is a no-op, so deny sites behave
145/// identically with metrics disabled. `limiter` is one of `tool`,
146/// `auth_pre`, `auth_post`, `extra_route`.
147pub(crate) fn record_rate_limit_deny(ext: &axum::http::Extensions, limiter: &str) {
148    if let Some(m) = ext.get::<Arc<McpMetrics>>() {
149        m.rate_limited_total.with_label_values(&[limiter]).inc();
150    }
151}
152
153/// Spawn a dedicated HTTP listener that serves Prometheus metrics on `/metrics`.
154///
155/// The listener exits and releases the bound port when `shutdown` is
156/// cancelled, keeping the metrics endpoint tied to the parent server's
157/// graceful-shutdown lifecycle (M7).
158///
159/// # Errors
160///
161/// Returns [`RmcpServerKitError::Startup`] if the TCP listener cannot bind or the
162/// underlying axum server fails.
163// cancel-safe: the parent server cancels via `shutdown.cancelled()` inside
164// axum graceful shutdown; dropping this future directly only drops the
165// listener/app, with no metrics registry mutation or detached work.
166pub async fn serve_metrics(
167    bind: String,
168    metrics: Arc<McpMetrics>,
169    shutdown: tokio_util::sync::CancellationToken,
170) -> Result<(), RmcpServerKitError> {
171    let app = axum::Router::new().route(
172        "/metrics",
173        axum::routing::get(move || {
174            let m = Arc::clone(&metrics);
175            async move { m.encode() }
176        }),
177    );
178
179    let listener = tokio::net::TcpListener::bind(&bind)
180        .await
181        .map_err(|e| RmcpServerKitError::Startup(format!("metrics bind {bind}: {e}")))?;
182    tracing::info!("metrics endpoint listening on http://{bind}/metrics");
183    axum::serve(listener, app)
184        .with_graceful_shutdown(async move { shutdown.cancelled().await })
185        .await
186        .map_err(|e| RmcpServerKitError::Startup(format!("metrics serve: {e}")))?;
187    Ok(())
188}
189
190#[cfg(test)]
191mod tests {
192    #![allow(
193        clippy::unwrap_used,
194        clippy::expect_used,
195        clippy::panic,
196        clippy::indexing_slicing,
197        clippy::unwrap_in_result,
198        clippy::print_stdout,
199        clippy::print_stderr,
200        reason = "test-only relaxations; production code uses ? and tracing"
201    )]
202    use super::*;
203
204    #[test]
205    fn encode_failure_returns_stable_non_empty_marker() {
206        // The encoder's only failure mode is an IO error from the sink, so the
207        // failure branch is driven here through `encode_failure_body` - the
208        // exact function `encode` returns through - rather than through a
209        // collector, which cannot make the encoder fail.
210        let body = encode_failure_body(&prometheus::Error::Msg("encoder exploded".to_owned()));
211
212        assert!(
213            !body.is_empty(),
214            "a failed encode must never serve an empty body"
215        );
216        assert!(
217            body.contains("rmcp-server-kit: prometheus encoding failed"),
218            "marker body must be stable for scrapers/alerts: {body:?}"
219        );
220        assert_eq!(body, ENCODE_FAILURE_MARKER);
221    }
222
223    #[test]
224    fn new_creates_registry_with_counters() {
225        let m = McpMetrics::new().unwrap();
226        // Incrementing a counter should make it appear in gather output.
227        m.http_requests_total
228            .with_label_values(&["GET", "/test", "200"])
229            .inc();
230        m.http_request_duration_seconds
231            .with_label_values(&["GET", "/test"])
232            .observe(0.1);
233        assert_eq!(m.registry.gather().len(), 2);
234    }
235
236    #[test]
237    fn encode_empty_registry() {
238        let m = McpMetrics::new().unwrap();
239        let output = m.encode();
240        // Empty counters/histograms produce no samples but the output is valid.
241        assert!(output.is_empty() || output.contains("rmcp_server_kit_"));
242    }
243
244    #[test]
245    fn counter_increment_shows_in_encode() {
246        let m = McpMetrics::new().unwrap();
247        m.http_requests_total
248            .with_label_values(&["GET", "/healthz", "200"])
249            .inc();
250        let output = m.encode();
251        assert!(output.contains("rmcp_server_kit_http_requests_total"));
252        assert!(output.contains("method=\"GET\""));
253        assert!(output.contains("path=\"/healthz\""));
254        assert!(output.contains("status=\"200\""));
255        assert!(output.contains(" 1")); // count = 1
256    }
257
258    #[test]
259    fn histogram_observe_shows_in_encode() {
260        let m = McpMetrics::new().unwrap();
261        m.http_request_duration_seconds
262            .with_label_values(&["POST", "/mcp"])
263            .observe(0.042);
264        let output = m.encode();
265        assert!(output.contains("rmcp_server_kit_http_request_duration_seconds"));
266        assert!(output.contains("method=\"POST\""));
267        assert!(output.contains("path=\"/mcp\""));
268    }
269
270    #[test]
271    fn multiple_increments_accumulate() {
272        let m = McpMetrics::new().unwrap();
273        let counter = m
274            .http_requests_total
275            .with_label_values(&["POST", "/mcp", "200"]);
276        counter.inc();
277        counter.inc();
278        counter.inc();
279        let output = m.encode();
280        assert!(output.contains(" 3")); // count = 3
281    }
282
283    #[test]
284    fn clone_shares_registry() {
285        let m = McpMetrics::new().unwrap();
286        let m2 = m.clone();
287        m.http_requests_total
288            .with_label_values(&["GET", "/test", "200"])
289            .inc();
290        // The clone should see the same counter value.
291        let output = m2.encode();
292        assert!(output.contains(" 1"));
293    }
294
295    #[test]
296    fn rate_limited_counter_registers_and_encodes() {
297        let m = McpMetrics::new().unwrap();
298        m.rate_limited_total.with_label_values(&["tool"]).inc();
299        let output = m.encode();
300        assert!(output.contains("rmcp_server_kit_rate_limited_total"));
301        assert!(output.contains("limiter=\"tool\""));
302        assert!(output.contains(" 1"));
303    }
304
305    #[test]
306    fn record_rate_limit_deny_increments_via_extension() {
307        let m = Arc::new(McpMetrics::new().unwrap());
308        let mut ext = axum::http::Extensions::new();
309        ext.insert(Arc::clone(&m));
310        record_rate_limit_deny(&ext, "auth_pre");
311        record_rate_limit_deny(&ext, "auth_pre");
312        assert_eq!(
313            m.rate_limited_total.with_label_values(&["auth_pre"]).get(),
314            2
315        );
316        // Absent handle: silent no-op (metrics disabled path).
317        let empty = axum::http::Extensions::new();
318        record_rate_limit_deny(&empty, "auth_pre");
319        assert_eq!(
320            m.rate_limited_total.with_label_values(&["auth_pre"]).get(),
321            2
322        );
323    }
324
325    // M7 regression: cancelling the shutdown token must release the
326    // metrics listener's bound port so a subsequent bind to the same
327    // address succeeds. Prior to M7 the metrics endpoint ran without
328    // graceful_shutdown wiring and would leak the port until process
329    // exit.
330    #[tokio::test]
331    async fn serve_metrics_releases_port_on_shutdown() {
332        // Pick an ephemeral port, then drop the probe so serve_metrics
333        // can claim it.
334        let probe = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
335        let addr = probe.local_addr().unwrap();
336        drop(probe);
337
338        let metrics = Arc::new(McpMetrics::new().unwrap());
339        let shutdown = tokio_util::sync::CancellationToken::new();
340        let handle = tokio::spawn(serve_metrics(
341            addr.to_string(),
342            Arc::clone(&metrics),
343            shutdown.clone(),
344        ));
345
346        // Wait until the listener is actually accepting connections.
347        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
348        loop {
349            if tokio::net::TcpStream::connect(addr).await.is_ok() {
350                break;
351            }
352            assert!(
353                std::time::Instant::now() < deadline,
354                "metrics listener never accepted on {addr}"
355            );
356            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
357        }
358
359        // Cancel and await graceful shutdown.
360        shutdown.cancel();
361        let join = tokio::time::timeout(std::time::Duration::from_secs(5), handle)
362            .await
363            .expect("serve_metrics did not return within timeout");
364        join.expect("join error")
365            .expect("serve_metrics returned Err");
366
367        // Port must be immediately rebindable.
368        let rebind = tokio::net::TcpListener::bind(addr)
369            .await
370            .expect("port not released after shutdown");
371        drop(rebind);
372    }
373}
rmcp_server_kit/metrics.rs

rmcp_server_kit/
metrics.rs
