agentmesh_mcp/mcp/
metrics.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Categorical metrics for MCP controls.
5
6use crate::mcp::error::McpError;
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9use std::sync::{Arc, Mutex};
10
11/// Decision labels recorded for `mcp_decisions`.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub enum McpDecisionLabel {
14    Allowed,
15    Denied,
16    RateLimited,
17    ApprovalRequired,
18    Sanitized,
19}
20
21impl McpDecisionLabel {
22    pub fn as_str(self) -> &'static str {
23        match self {
24            Self::Allowed => "allowed",
25            Self::Denied => "denied",
26            Self::RateLimited => "rate_limited",
27            Self::ApprovalRequired => "approval_required",
28            Self::Sanitized => "sanitized",
29        }
30    }
31}
32
33/// Threat labels recorded for `mcp_threats_detected`.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
35pub enum McpThreatLabel {
36    ToolPoisoning,
37    RugPull,
38    CrossServerAttack,
39    DescriptionInjection,
40    SchemaAbuse,
41    HiddenInstruction,
42    PromptInjectionTag,
43    ImperativePhrasing,
44    CredentialLeakage,
45    ExfiltrationUrl,
46}
47
48impl McpThreatLabel {
49    pub fn as_str(self) -> &'static str {
50        match self {
51            Self::ToolPoisoning => "tool_poisoning",
52            Self::RugPull => "rug_pull",
53            Self::CrossServerAttack => "cross_server_attack",
54            Self::DescriptionInjection => "description_injection",
55            Self::SchemaAbuse => "schema_abuse",
56            Self::HiddenInstruction => "hidden_instruction",
57            Self::PromptInjectionTag => "prompt_injection_tag",
58            Self::ImperativePhrasing => "imperative_phrasing",
59            Self::CredentialLeakage => "credential_leakage",
60            Self::ExfiltrationUrl => "exfiltration_url",
61        }
62    }
63}
64
65/// Scan labels recorded for `mcp_scans`.
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
67pub enum McpScanLabel {
68    Response,
69    ToolMetadata,
70    Gateway,
71}
72
73impl McpScanLabel {
74    pub fn as_str(self) -> &'static str {
75        match self {
76            Self::Response => "response",
77            Self::ToolMetadata => "tool_metadata",
78            Self::Gateway => "gateway",
79        }
80    }
81}
82
83#[derive(Debug, Default)]
84struct McpMetricsState {
85    decisions: BTreeMap<String, u64>,
86    threats_detected: BTreeMap<String, u64>,
87    rate_limit_hits: BTreeMap<String, u64>,
88    scans: BTreeMap<String, u64>,
89}
90
91/// Thread-safe categorical metrics collector.
92#[derive(Debug, Clone, Default)]
93pub struct McpMetricsCollector {
94    state: Arc<Mutex<McpMetricsState>>,
95}
96
97/// Snapshot of collected categorical metrics.
98#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
99pub struct McpMetricsSnapshot {
100    pub mcp_decisions: BTreeMap<String, u64>,
101    pub mcp_threats_detected: BTreeMap<String, u64>,
102    pub mcp_rate_limit_hits: BTreeMap<String, u64>,
103    pub mcp_scans: BTreeMap<String, u64>,
104}
105
106impl McpMetricsCollector {
107    pub fn record_decision(&self, label: McpDecisionLabel) -> Result<(), McpError> {
108        self.increment("metrics", |state| &mut state.decisions, label.as_str())
109    }
110
111    pub fn record_threat(&self, label: McpThreatLabel) -> Result<(), McpError> {
112        self.increment(
113            "metrics",
114            |state| &mut state.threats_detected,
115            label.as_str(),
116        )
117    }
118
119    pub fn record_rate_limit_hit(&self, label: &str) -> Result<(), McpError> {
120        self.increment("metrics", |state| &mut state.rate_limit_hits, label)
121    }
122
123    pub fn record_scan(&self, label: McpScanLabel) -> Result<(), McpError> {
124        self.increment("metrics", |state| &mut state.scans, label.as_str())
125    }
126
127    pub fn snapshot(&self) -> Result<McpMetricsSnapshot, McpError> {
128        let state = self
129            .state
130            .lock()
131            .map_err(|_| McpError::store("metrics", "metrics lock poisoned"))?;
132        Ok(McpMetricsSnapshot {
133            mcp_decisions: state.decisions.clone(),
134            mcp_threats_detected: state.threats_detected.clone(),
135            mcp_rate_limit_hits: state.rate_limit_hits.clone(),
136            mcp_scans: state.scans.clone(),
137        })
138    }
139
140    fn increment<F>(&self, store: &'static str, select: F, key: &str) -> Result<(), McpError>
141    where
142        F: FnOnce(&mut McpMetricsState) -> &mut BTreeMap<String, u64>,
143    {
144        let mut state = self
145            .state
146            .lock()
147            .map_err(|_| McpError::store(store, "metrics lock poisoned"))?;
148        let entry = select(&mut state).entry(key.to_string()).or_insert(0);
149        *entry += 1;
150        Ok(())
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn metrics_snapshot_is_categorical() {
160        let metrics = McpMetricsCollector::default();
161        metrics.record_decision(McpDecisionLabel::Allowed).unwrap();
162        metrics.record_threat(McpThreatLabel::SchemaAbuse).unwrap();
163        let snapshot = metrics.snapshot().unwrap();
164        assert_eq!(snapshot.mcp_decisions["allowed"], 1);
165        assert_eq!(snapshot.mcp_threats_detected["schema_abuse"], 1);
166    }
167}
agentmesh_mcp/mcp/metrics.rs

agentmesh_mcp/mcp/
metrics.rs
