mz_environmentd/http/
mcp_metrics.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Prometheus metrics for the MCP HTTP endpoints.
11//!
12//! Tracks request counts, tool call counts, and tool call durations,
13//! labeled by endpoint type (`agent` / `developer`) and either the
14//! JSON-RPC method name or the MCP tool name. The status label is one of
15//! the [`McpCallStatus`] variants: `ok`, a lifecycle outcome, or the
16//! `McpRequestError` error type (e.g. `ToolNotFound`) for failures.
17
18use mz_ore::metric;
19use mz_ore::metrics::MetricsRegistry;
20use mz_ore::stats::histogram_seconds_buckets;
21use prometheus::{HistogramTimer, HistogramVec, IntCounterVec};
22
23/// Closed set of outcomes recorded in the MCP `status` label. Keeping these
24/// as an enum (rather than free-form strings at the call sites) pins the
25/// metric's label cardinality and stops typos from silently creating new
26/// label values. Mirrors `oauth_metadata::MetricStatus`.
27#[derive(Debug, Clone, Copy)]
28pub enum McpCallStatus {
29    /// The request or tool call succeeded.
30    Ok,
31    /// The surrounding future was dropped before the outcome was recorded
32    /// (e.g. by the request timeout). This is the [`ToolCallGuard`] default,
33    /// so a dropped call records as cancelled rather than being lost.
34    Cancelled,
35    /// The request exceeded the configured timeout.
36    Timeout,
37    /// The endpoint is disabled by its feature flag.
38    EndpointDisabled,
39    /// The request failed with the given `McpRequestError::error_type()`. The
40    /// wrapped value comes from a closed match, so cardinality stays bounded.
41    Error(&'static str),
42}
43
44impl McpCallStatus {
45    pub fn as_str(self) -> &'static str {
46        match self {
47            Self::Ok => "ok",
48            Self::Cancelled => "cancelled",
49            Self::Timeout => "timeout",
50            Self::EndpointDisabled => "endpoint_disabled",
51            Self::Error(e) => e,
52        }
53    }
54}
55
56/// Metrics emitted by the MCP HTTP handlers.
57///
58/// Cheaply `Clone`: Prometheus collector handles are `Arc`-shared internally,
59/// so the struct can be cloned freely and stored as an axum `Extension`.
60#[derive(Debug, Clone)]
61pub struct McpMetrics {
62    /// Total MCP requests by endpoint type, JSON-RPC method, and status.
63    pub requests: IntCounterVec,
64    /// Total MCP `tools/call` invocations by endpoint type, tool name, and status.
65    pub tool_calls: IntCounterVec,
66    /// Duration of MCP `tools/call` invocations by endpoint type and tool name.
67    pub tool_call_duration: HistogramVec,
68}
69
70/// RAII guard for a single `tools/call` invocation. On drop, increments
71/// `tool_calls_total` with the current status and observes
72/// `tool_call_duration_seconds` via the embedded [`HistogramTimer`]'s own
73/// drop. Designed so that if the surrounding future is dropped before
74/// completion (e.g. by `tokio::time::timeout`), the metric still records
75/// with the default `"cancelled"` status instead of being silently lost.
76pub struct ToolCallGuard<'a> {
77    metrics: &'a McpMetrics,
78    endpoint_label: &'static str,
79    tool_label: String,
80    status: McpCallStatus,
81    /// `HistogramTimer::drop` observes the duration into the histogram, so
82    /// holding the timer here means we get the duration recorded for both
83    /// normal completion and early drop.
84    _timer: HistogramTimer,
85}
86
87impl<'a> ToolCallGuard<'a> {
88    /// Starts a new tool call: begins the duration timer and reserves the
89    /// counter increment that will happen on drop.
90    pub fn new(metrics: &'a McpMetrics, endpoint_label: &'static str, tool_label: String) -> Self {
91        let timer = metrics
92            .tool_call_duration
93            .with_label_values(&[endpoint_label, &tool_label])
94            .start_timer();
95        Self {
96            metrics,
97            endpoint_label,
98            tool_label,
99            status: McpCallStatus::Cancelled,
100            _timer: timer,
101        }
102    }
103
104    /// Records the outcome of the call. Callers should set this on the
105    /// normal completion path right before the guard is dropped.
106    pub fn set_status(&mut self, status: McpCallStatus) {
107        self.status = status;
108    }
109}
110
111impl Drop for ToolCallGuard<'_> {
112    fn drop(&mut self) {
113        self.metrics
114            .tool_calls
115            .with_label_values(&[self.endpoint_label, &self.tool_label, self.status.as_str()])
116            .inc();
117    }
118}
119
120impl McpMetrics {
121    pub fn register_into(registry: &MetricsRegistry) -> Self {
122        Self {
123            requests: registry.register(metric!(
124                name: "mz_mcp_requests_total",
125                help: "Total number of MCP requests received.",
126                var_labels: ["endpoint_type", "method", "status"],
127            )),
128            tool_calls: registry.register(metric!(
129                name: "mz_mcp_tool_calls_total",
130                help: "Total number of MCP tools/call invocations.",
131                var_labels: ["endpoint_type", "tool_name", "status"],
132            )),
133            tool_call_duration: registry.register(metric!(
134                name: "mz_mcp_tool_call_duration_seconds",
135                help: "Duration of MCP tools/call invocations in seconds.",
136                var_labels: ["endpoint_type", "tool_name"],
137                buckets: histogram_seconds_buckets(0.000_128, 8.0),
138            )),
139        }
140    }
141
142    /// Records a request in `mz_mcp_requests_total`. Encapsulates the label
143    /// ordering and the [`McpCallStatus`] to `&str` conversion so call sites
144    /// never touch raw label strings.
145    pub fn record_request(&self, endpoint_label: &str, method_label: &str, status: McpCallStatus) {
146        self.requests
147            .with_label_values(&[endpoint_label, method_label, status.as_str()])
148            .inc();
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::{McpCallStatus, McpMetrics};
155    use mz_ore::metrics::MetricsRegistry;
156
157    /// The status label strings are a wire/dashboard contract; pin them so a
158    /// rename is a deliberate, visible change.
159    #[mz_ore::test]
160    fn test_call_status_labels() {
161        assert_eq!(McpCallStatus::Ok.as_str(), "ok");
162        assert_eq!(McpCallStatus::Cancelled.as_str(), "cancelled");
163        assert_eq!(McpCallStatus::Timeout.as_str(), "timeout");
164        assert_eq!(
165            McpCallStatus::EndpointDisabled.as_str(),
166            "endpoint_disabled"
167        );
168        assert_eq!(
169            McpCallStatus::Error("ToolNotFound").as_str(),
170            "ToolNotFound"
171        );
172    }
173
174    /// All three metrics register cleanly and show up in the gathered output
175    /// with the expected names. `IntCounterVec` / `HistogramVec` families
176    /// only appear in `gather()` after at least one label combination has
177    /// been observed, so each metric is touched once before gathering.
178    #[mz_ore::test]
179    fn test_register_into() {
180        let registry = MetricsRegistry::new();
181        let metrics = McpMetrics::register_into(&registry);
182
183        metrics
184            .requests
185            .with_label_values(&["agent", "initialize", "ok"])
186            .inc_by(0);
187        metrics
188            .tool_calls
189            .with_label_values(&["agent", "read_data_product", "ok"])
190            .inc_by(0);
191        metrics
192            .tool_call_duration
193            .with_label_values(&["agent", "read_data_product"])
194            .observe(0.0);
195
196        let names: Vec<String> = registry
197            .gather()
198            .iter()
199            .map(|m| m.name().to_string())
200            .collect();
201
202        assert!(
203            names.iter().any(|n| n == "mz_mcp_requests_total"),
204            "mz_mcp_requests_total should be registered, got: {names:?}",
205        );
206        assert!(
207            names.iter().any(|n| n == "mz_mcp_tool_calls_total"),
208            "mz_mcp_tool_calls_total should be registered, got: {names:?}",
209        );
210        assert!(
211            names
212                .iter()
213                .any(|n| n == "mz_mcp_tool_call_duration_seconds"),
214            "mz_mcp_tool_call_duration_seconds should be registered, got: {names:?}",
215        );
216    }
217
218    /// Incrementing each counter with realistic label values produces the
219    /// expected counts in the gathered output.
220    #[mz_ore::test]
221    fn test_record_metrics() {
222        let registry = MetricsRegistry::new();
223        let metrics = McpMetrics::register_into(&registry);
224
225        metrics
226            .requests
227            .with_label_values(&["agent", "tools/call", "ok"])
228            .inc();
229        metrics
230            .requests
231            .with_label_values(&["agent", "tools/call", "ok"])
232            .inc();
233        metrics
234            .requests
235            .with_label_values(&["developer", "initialize", "ok"])
236            .inc();
237
238        metrics
239            .tool_calls
240            .with_label_values(&["agent", "read_data_product", "ok"])
241            .inc();
242        metrics
243            .tool_calls
244            .with_label_values(&["agent", "read_data_product", "DataProductNotFound"])
245            .inc();
246
247        metrics
248            .tool_call_duration
249            .with_label_values(&["agent", "read_data_product"])
250            .observe(0.123);
251
252        let gathered = registry.gather();
253
254        // requests_total: 3 increments produce 2 distinct label sets (the
255        // first two share labels and so collapse into the same series).
256        let requests = gathered
257            .iter()
258            .find(|m| m.name() == "mz_mcp_requests_total")
259            .expect("requests metric present");
260        assert_eq!(requests.get_metric().len(), 2);
261
262        // tool_calls_total: 2 distinct label sets (one for each status).
263        let tool_calls = gathered
264            .iter()
265            .find(|m| m.name() == "mz_mcp_tool_calls_total")
266            .expect("tool_calls metric present");
267        assert_eq!(tool_calls.get_metric().len(), 2);
268
269        // tool_call_duration_seconds: one observation in one bucket set.
270        let duration = gathered
271            .iter()
272            .find(|m| m.name() == "mz_mcp_tool_call_duration_seconds")
273            .expect("tool_call_duration metric present");
274        assert_eq!(
275            duration.get_metric()[0].get_histogram().get_sample_count(),
276            1
277        );
278    }
279}
mz_environmentd/http/mcp_metrics.rs

mz_environmentd/http/
mcp_metrics.rs
