Skip to main content

mcpkit_core/debug/
inspector.rs

1//! Message inspection utilities.
2//!
3//! The message inspector captures protocol messages for analysis
4//! and debugging.
5
6use crate::protocol::{Message, RequestId};
7use std::collections::HashMap;
8use std::sync::{Arc, RwLock};
9use std::time::{Duration, Instant};
10
11/// A record of a captured message.
12#[derive(Debug, Clone)]
13pub struct MessageRecord {
14    /// Timestamp when the message was captured.
15    pub timestamp: Instant,
16    /// The direction of the message.
17    pub direction: MessageDirection,
18    /// The captured message.
19    pub message: Message,
20    /// Size in bytes (approximate).
21    pub size_bytes: usize,
22    /// Optional context/tags.
23    pub tags: Vec<String>,
24}
25
26impl MessageRecord {
27    /// Create a new message record.
28    #[must_use]
29    pub fn new(direction: MessageDirection, message: Message) -> Self {
30        let size_bytes = serde_json::to_string(&message).map_or(0, |s| s.len());
31
32        Self {
33            timestamp: Instant::now(),
34            direction,
35            message,
36            size_bytes,
37            tags: Vec::new(),
38        }
39    }
40
41    /// Add a tag to the record.
42    #[must_use]
43    pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
44        self.tags.push(tag.into());
45        self
46    }
47
48    /// Get the method name if applicable.
49    #[must_use]
50    pub fn method(&self) -> Option<&str> {
51        self.message.method()
52    }
53
54    /// Get the request ID if applicable.
55    #[must_use]
56    pub fn request_id(&self) -> Option<&RequestId> {
57        match &self.message {
58            Message::Request(req) => Some(&req.id),
59            Message::Response(res) => Some(&res.id),
60            Message::Notification(_) => None,
61        }
62    }
63
64    /// Check if this is a request message.
65    #[must_use]
66    pub fn is_request(&self) -> bool {
67        matches!(self.message, Message::Request(_))
68    }
69
70    /// Check if this is a response message.
71    #[must_use]
72    pub fn is_response(&self) -> bool {
73        matches!(self.message, Message::Response(_))
74    }
75
76    /// Check if this is a notification message.
77    #[must_use]
78    pub fn is_notification(&self) -> bool {
79        matches!(self.message, Message::Notification(_))
80    }
81
82    /// Check if the response indicates an error.
83    #[must_use]
84    pub fn is_error(&self) -> bool {
85        matches!(&self.message, Message::Response(r) if r.error.is_some())
86    }
87}
88
89/// Direction of a message.
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum MessageDirection {
92    /// Message sent by client.
93    Outbound,
94    /// Message received by client.
95    Inbound,
96}
97
98/// Statistics about captured messages.
99#[derive(Debug, Clone, Default)]
100pub struct MessageStats {
101    /// Total messages captured.
102    pub total_messages: usize,
103    /// Total requests.
104    pub requests: usize,
105    /// Total responses.
106    pub responses: usize,
107    /// Total notifications.
108    pub notifications: usize,
109    /// Total errors (error responses).
110    pub errors: usize,
111    /// Total bytes transferred.
112    pub total_bytes: usize,
113    /// Messages by method.
114    pub by_method: HashMap<String, usize>,
115    /// Average response time per method.
116    pub avg_response_time: HashMap<String, Duration>,
117}
118
119impl MessageStats {
120    /// Calculate error rate.
121    #[must_use]
122    pub fn error_rate(&self) -> f64 {
123        if self.responses == 0 {
124            0.0
125        } else {
126            self.errors as f64 / self.responses as f64
127        }
128    }
129}
130
131/// Message inspector for capturing and analyzing protocol traffic.
132///
133/// The inspector can be used as a debugging aid during development
134/// to understand the message flow between client and server.
135#[derive(Debug)]
136pub struct MessageInspector {
137    /// Captured messages.
138    records: Arc<RwLock<Vec<MessageRecord>>>,
139    /// Maximum number of records to keep (0 = unlimited).
140    max_records: usize,
141    /// Whether capturing is enabled.
142    enabled: Arc<RwLock<bool>>,
143    /// Pending requests for response time tracking.
144    pending_requests: Arc<RwLock<HashMap<RequestId, (Instant, String)>>>,
145    /// Response times by method.
146    response_times: Arc<RwLock<HashMap<String, Vec<Duration>>>>,
147}
148
149impl Default for MessageInspector {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl MessageInspector {
156    /// Create a new message inspector.
157    #[must_use]
158    pub fn new() -> Self {
159        Self {
160            records: Arc::new(RwLock::new(Vec::new())),
161            max_records: 10000,
162            enabled: Arc::new(RwLock::new(true)),
163            pending_requests: Arc::new(RwLock::new(HashMap::new())),
164            response_times: Arc::new(RwLock::new(HashMap::new())),
165        }
166    }
167
168    /// Create an inspector with a maximum record limit.
169    #[must_use]
170    pub fn with_max_records(mut self, max: usize) -> Self {
171        self.max_records = max;
172        self
173    }
174
175    /// Enable or disable message capture.
176    pub fn set_enabled(&self, enabled: bool) {
177        if let Ok(mut flag) = self.enabled.write() {
178            *flag = enabled;
179        }
180    }
181
182    /// Check if capturing is enabled.
183    #[must_use]
184    pub fn is_enabled(&self) -> bool {
185        self.enabled.read().is_ok_and(|e| *e)
186    }
187
188    /// Record an outbound message.
189    pub fn record_outbound(&self, message: Message) {
190        self.record(MessageDirection::Outbound, message);
191    }
192
193    /// Record an inbound message.
194    pub fn record_inbound(&self, message: Message) {
195        self.record(MessageDirection::Inbound, message);
196    }
197
198    /// Record a message.
199    fn record(&self, direction: MessageDirection, message: Message) {
200        if !self.is_enabled() {
201            return;
202        }
203
204        let record = MessageRecord::new(direction, message.clone());
205
206        // Track pending requests for response time calculation
207        if let Message::Request(ref req) = message {
208            if let Ok(mut pending) = self.pending_requests.write() {
209                pending.insert(req.id.clone(), (Instant::now(), req.method.to_string()));
210            }
211        }
212
213        // Calculate response time for completed requests
214        if let Message::Response(ref res) = message {
215            if let Ok(mut pending) = self.pending_requests.write() {
216                if let Some((start, method)) = pending.remove(&res.id) {
217                    let duration = start.elapsed();
218                    if let Ok(mut times) = self.response_times.write() {
219                        times.entry(method).or_default().push(duration);
220                    }
221                }
222            }
223        }
224
225        // Store the record
226        if let Ok(mut records) = self.records.write() {
227            records.push(record);
228
229            // Trim if over limit
230            if self.max_records > 0 && records.len() > self.max_records {
231                let excess = records.len() - self.max_records;
232                records.drain(0..excess);
233            }
234        }
235    }
236
237    /// Get all captured records.
238    #[must_use]
239    pub fn records(&self) -> Vec<MessageRecord> {
240        self.records.read().map(|r| r.clone()).unwrap_or_default()
241    }
242
243    /// Get the number of captured records.
244    #[must_use]
245    pub fn len(&self) -> usize {
246        self.records.read().map_or(0, |r| r.len())
247    }
248
249    /// Check if there are no captured records.
250    #[must_use]
251    pub fn is_empty(&self) -> bool {
252        self.len() == 0
253    }
254
255    /// Clear all captured records.
256    pub fn clear(&self) {
257        if let Ok(mut records) = self.records.write() {
258            records.clear();
259        }
260        if let Ok(mut pending) = self.pending_requests.write() {
261            pending.clear();
262        }
263        if let Ok(mut times) = self.response_times.write() {
264            times.clear();
265        }
266    }
267
268    /// Get message statistics.
269    #[must_use]
270    #[allow(clippy::field_reassign_with_default)]
271    pub fn stats(&self) -> MessageStats {
272        let records = self.records();
273        let mut stats = MessageStats::default();
274
275        stats.total_messages = records.len();
276
277        for record in &records {
278            stats.total_bytes += record.size_bytes;
279
280            match &record.message {
281                Message::Request(req) => {
282                    stats.requests += 1;
283                    *stats.by_method.entry(req.method.to_string()).or_insert(0) += 1;
284                }
285                Message::Response(res) => {
286                    stats.responses += 1;
287                    if res.error.is_some() {
288                        stats.errors += 1;
289                    }
290                }
291                Message::Notification(notif) => {
292                    stats.notifications += 1;
293                    *stats.by_method.entry(notif.method.to_string()).or_insert(0) += 1;
294                }
295            }
296        }
297
298        // Calculate average response times
299        if let Ok(times) = self.response_times.read() {
300            for (method, durations) in times.iter() {
301                if !durations.is_empty() {
302                    let total: Duration = durations.iter().sum();
303                    let avg = total / durations.len() as u32;
304                    stats.avg_response_time.insert(method.clone(), avg);
305                }
306            }
307        }
308
309        stats
310    }
311
312    /// Find requests without responses.
313    #[must_use]
314    pub fn pending_requests(&self) -> Vec<(RequestId, String, Duration)> {
315        self.pending_requests
316            .read()
317            .map(|pending| {
318                pending
319                    .iter()
320                    .map(|(id, (start, method))| (id.clone(), method.clone(), start.elapsed()))
321                    .collect()
322            })
323            .unwrap_or_default()
324    }
325
326    /// Filter records by method.
327    #[must_use]
328    pub fn filter_by_method(&self, method: &str) -> Vec<MessageRecord> {
329        self.records()
330            .into_iter()
331            .filter(|r| r.method() == Some(method))
332            .collect()
333    }
334
335    /// Filter records by direction.
336    #[must_use]
337    pub fn filter_by_direction(&self, direction: MessageDirection) -> Vec<MessageRecord> {
338        self.records()
339            .into_iter()
340            .filter(|r| r.direction == direction)
341            .collect()
342    }
343
344    /// Get only error responses.
345    #[must_use]
346    pub fn errors(&self) -> Vec<MessageRecord> {
347        self.records()
348            .into_iter()
349            .filter(MessageRecord::is_error)
350            .collect()
351    }
352
353    /// Export records to JSON.
354    ///
355    /// # Errors
356    ///
357    /// Returns an error if serialization fails.
358    pub fn to_json(&self) -> Result<String, serde_json::Error> {
359        let records: Vec<_> = self
360            .records()
361            .into_iter()
362            .map(|r| {
363                serde_json::json!({
364                    "direction": format!("{:?}", r.direction),
365                    "message": r.message,
366                    "size_bytes": r.size_bytes,
367                    "tags": r.tags,
368                })
369            })
370            .collect();
371
372        serde_json::to_string_pretty(&records)
373    }
374}
375
376impl Clone for MessageInspector {
377    fn clone(&self) -> Self {
378        Self {
379            records: Arc::clone(&self.records),
380            max_records: self.max_records,
381            enabled: Arc::clone(&self.enabled),
382            pending_requests: Arc::clone(&self.pending_requests),
383            response_times: Arc::clone(&self.response_times),
384        }
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use crate::protocol::{Request, Response};
392
393    #[test]
394    fn test_message_inspector() {
395        let inspector = MessageInspector::new();
396
397        // Record some messages
398        let req = Message::Request(Request::new("test/method", 1));
399        inspector.record_outbound(req);
400
401        let res = Message::Response(Response::success(RequestId::from(1), serde_json::json!({})));
402        inspector.record_inbound(res);
403
404        assert_eq!(inspector.len(), 2);
405
406        let stats = inspector.stats();
407        assert_eq!(stats.requests, 1);
408        assert_eq!(stats.responses, 1);
409        assert_eq!(stats.errors, 0);
410    }
411
412    #[test]
413    fn test_filter_by_method() {
414        let inspector = MessageInspector::new();
415
416        inspector.record_outbound(Message::Request(Request::new("method/a", 1)));
417        inspector.record_outbound(Message::Request(Request::new("method/b", 2)));
418        inspector.record_outbound(Message::Request(Request::new("method/a", 3)));
419
420        let filtered = inspector.filter_by_method("method/a");
421        assert_eq!(filtered.len(), 2);
422    }
423
424    #[test]
425    fn test_max_records() {
426        let inspector = MessageInspector::new().with_max_records(5);
427
428        for i in 0..10 {
429            inspector.record_outbound(Message::Request(Request::new("test", i)));
430        }
431
432        assert_eq!(inspector.len(), 5);
433    }
434
435    #[test]
436    fn test_enable_disable() {
437        let inspector = MessageInspector::new();
438
439        inspector.record_outbound(Message::Request(Request::new("test", 1)));
440        assert_eq!(inspector.len(), 1);
441
442        inspector.set_enabled(false);
443        inspector.record_outbound(Message::Request(Request::new("test", 2)));
444        assert_eq!(inspector.len(), 1); // Still 1, not recorded
445    }
446}