use super::types::{McpOperationDetails, RequestMetadata, TraceContext};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpRequestEvent {
pub trace: TraceContext,
pub server_name: String,
pub operation: McpOperationDetails,
pub metadata: RequestMetadata,
#[serde(skip_serializing_if = "Option::is_none")]
pub user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tenant_id: Option<String>,
pub timestamp: DateTime<Utc>,
}
impl McpRequestEvent {
pub fn new(
trace: TraceContext,
server_name: impl Into<String>,
operation: McpOperationDetails,
) -> Self {
Self {
trace,
server_name: server_name.into(),
operation,
metadata: RequestMetadata::default(),
user_id: None,
tenant_id: None,
timestamp: Utc::now(),
}
}
pub fn with_metadata(mut self, metadata: RequestMetadata) -> Self {
self.metadata = metadata;
self
}
pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
self.user_id = Some(user_id.into());
self
}
pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
self.tenant_id = Some(tenant_id.into());
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpResponseEvent {
pub trace: TraceContext,
pub server_name: String,
pub operation: McpOperationDetails,
pub metadata: RequestMetadata,
#[serde(skip_serializing_if = "Option::is_none")]
pub user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tenant_id: Option<String>,
pub duration_ms: u64,
pub success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_code: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_message: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub response_size: Option<usize>,
pub timestamp: DateTime<Utc>,
}
impl McpResponseEvent {
pub fn success(
trace: TraceContext,
server_name: impl Into<String>,
operation: McpOperationDetails,
duration_ms: u64,
) -> Self {
Self {
trace,
server_name: server_name.into(),
operation,
metadata: RequestMetadata::default(),
user_id: None,
tenant_id: None,
duration_ms,
success: true,
error_code: None,
error_message: None,
response_size: None,
timestamp: Utc::now(),
}
}
pub fn failure(
trace: TraceContext,
server_name: impl Into<String>,
operation: McpOperationDetails,
duration_ms: u64,
error_code: i32,
error_message: impl Into<String>,
) -> Self {
Self {
trace,
server_name: server_name.into(),
operation,
metadata: RequestMetadata::default(),
user_id: None,
tenant_id: None,
duration_ms,
success: false,
error_code: Some(error_code),
error_message: Some(error_message.into()),
response_size: None,
timestamp: Utc::now(),
}
}
pub fn with_metadata(mut self, metadata: RequestMetadata) -> Self {
self.metadata = metadata;
self
}
pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
self.user_id = Some(user_id.into());
self
}
pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
self.tenant_id = Some(tenant_id.into());
self
}
pub fn with_response_size(mut self, size: usize) -> Self {
self.response_size = Some(size);
self
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum MetricUnit {
Milliseconds,
Count,
Bytes,
Percent,
None,
}
impl MetricUnit {
pub fn as_str(&self) -> &'static str {
match self {
Self::Milliseconds => "Milliseconds",
Self::Count => "Count",
Self::Bytes => "Bytes",
Self::Percent => "Percent",
Self::None => "None",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpMetric {
pub name: String,
pub value: f64,
pub unit: MetricUnit,
pub dimensions: HashMap<String, String>,
pub timestamp: DateTime<Utc>,
}
impl McpMetric {
pub fn new(name: impl Into<String>, value: f64, unit: MetricUnit) -> Self {
Self {
name: name.into(),
value,
unit,
dimensions: HashMap::new(),
timestamp: Utc::now(),
}
}
pub fn with_dimension(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.dimensions.insert(key.into(), value.into());
self
}
pub fn with_dimensions(mut self, dimensions: HashMap<String, String>) -> Self {
self.dimensions.extend(dimensions);
self
}
pub fn duration(name: impl Into<String>, duration_ms: u64) -> Self {
Self::new(name, duration_ms as f64, MetricUnit::Milliseconds)
}
pub fn count(name: impl Into<String>, count: u64) -> Self {
Self::new(name, count as f64, MetricUnit::Count)
}
pub fn bytes(name: impl Into<String>, bytes: usize) -> Self {
Self::new(name, bytes as f64, MetricUnit::Bytes)
}
}
#[derive(Debug, Clone, Copy)]
pub struct StandardMetrics;
impl StandardMetrics {
pub const REQUEST_DURATION: &'static str = "mcp.request.duration";
pub const REQUEST_COUNT: &'static str = "mcp.request.count";
pub const REQUEST_ERRORS: &'static str = "mcp.request.errors";
pub const RESPONSE_SIZE: &'static str = "mcp.response.size";
pub const COMPOSITION_DEPTH: &'static str = "mcp.composition.depth";
}
#[derive(Debug, Clone)]
pub struct RequestStart {
pub instant: Instant,
pub trace: TraceContext,
pub operation: McpOperationDetails,
pub metadata: RequestMetadata,
}
impl RequestStart {
pub fn new(trace: TraceContext, operation: McpOperationDetails) -> Self {
Self {
instant: Instant::now(),
trace,
operation,
metadata: RequestMetadata::default(),
}
}
pub fn elapsed_ms(&self) -> u64 {
self.instant.elapsed().as_millis() as u64
}
pub fn with_metadata(mut self, metadata: RequestMetadata) -> Self {
self.metadata = metadata;
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_request_event_creation() {
let trace = TraceContext::new_root();
let operation = McpOperationDetails::tool_call("get_weather");
let event = McpRequestEvent::new(trace.clone(), "test-server", operation)
.with_user_id("user-123")
.with_tenant_id("tenant-456");
assert_eq!(event.server_name, "test-server");
assert_eq!(event.user_id, Some("user-123".to_string()));
assert_eq!(event.tenant_id, Some("tenant-456".to_string()));
assert_eq!(event.trace.trace_id, trace.trace_id);
}
#[test]
fn test_response_event_success() {
let trace = TraceContext::new_root();
let operation = McpOperationDetails::tool_call("get_weather");
let event = McpResponseEvent::success(trace.clone(), "test-server", operation, 150)
.with_response_size(1024);
assert!(event.success);
assert_eq!(event.duration_ms, 150);
assert_eq!(event.response_size, Some(1024));
assert!(event.error_code.is_none());
}
#[test]
fn test_response_event_failure() {
let trace = TraceContext::new_root();
let operation = McpOperationDetails::tool_call("get_weather");
let event = McpResponseEvent::failure(
trace.clone(),
"test-server",
operation,
50,
-32600,
"Invalid request",
);
assert!(!event.success);
assert_eq!(event.error_code, Some(-32600));
assert_eq!(event.error_message, Some("Invalid request".to_string()));
}
#[test]
fn test_metric_creation() {
let metric = McpMetric::duration("mcp.request.duration", 150)
.with_dimension("server", "test-server")
.with_dimension("method", "tools/call");
assert_eq!(metric.name, "mcp.request.duration");
assert!((metric.value - 150.0).abs() < f64::EPSILON);
assert_eq!(metric.unit, MetricUnit::Milliseconds);
assert_eq!(
metric.dimensions.get("server"),
Some(&"test-server".to_string())
);
}
#[test]
fn test_request_start_elapsed() {
let trace = TraceContext::new_root();
let operation = McpOperationDetails::tool_call("test");
let start = RequestStart::new(trace, operation);
let elapsed = start.elapsed_ms();
assert!(elapsed < 1000); }
#[test]
fn test_metric_unit_as_str() {
assert_eq!(MetricUnit::Milliseconds.as_str(), "Milliseconds");
assert_eq!(MetricUnit::Count.as_str(), "Count");
assert_eq!(MetricUnit::Bytes.as_str(), "Bytes");
assert_eq!(MetricUnit::Percent.as_str(), "Percent");
assert_eq!(MetricUnit::None.as_str(), "None");
}
}