1use crate::protocol::{Message, RequestId};
7use std::collections::HashMap;
8use std::sync::{Arc, RwLock};
9use std::time::{Duration, Instant};
10
11#[derive(Debug, Clone)]
13pub struct MessageRecord {
14 pub timestamp: Instant,
16 pub direction: MessageDirection,
18 pub message: Message,
20 pub size_bytes: usize,
22 pub tags: Vec<String>,
24}
25
26impl MessageRecord {
27 #[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 #[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 #[must_use]
50 pub fn method(&self) -> Option<&str> {
51 self.message.method()
52 }
53
54 #[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 #[must_use]
66 pub fn is_request(&self) -> bool {
67 matches!(self.message, Message::Request(_))
68 }
69
70 #[must_use]
72 pub fn is_response(&self) -> bool {
73 matches!(self.message, Message::Response(_))
74 }
75
76 #[must_use]
78 pub fn is_notification(&self) -> bool {
79 matches!(self.message, Message::Notification(_))
80 }
81
82 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum MessageDirection {
92 Outbound,
94 Inbound,
96}
97
98#[derive(Debug, Clone, Default)]
100pub struct MessageStats {
101 pub total_messages: usize,
103 pub requests: usize,
105 pub responses: usize,
107 pub notifications: usize,
109 pub errors: usize,
111 pub total_bytes: usize,
113 pub by_method: HashMap<String, usize>,
115 pub avg_response_time: HashMap<String, Duration>,
117}
118
119impl MessageStats {
120 #[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#[derive(Debug)]
136pub struct MessageInspector {
137 records: Arc<RwLock<Vec<MessageRecord>>>,
139 max_records: usize,
141 enabled: Arc<RwLock<bool>>,
143 pending_requests: Arc<RwLock<HashMap<RequestId, (Instant, String)>>>,
145 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 #[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 #[must_use]
170 pub fn with_max_records(mut self, max: usize) -> Self {
171 self.max_records = max;
172 self
173 }
174
175 pub fn set_enabled(&self, enabled: bool) {
177 if let Ok(mut flag) = self.enabled.write() {
178 *flag = enabled;
179 }
180 }
181
182 #[must_use]
184 pub fn is_enabled(&self) -> bool {
185 self.enabled.read().is_ok_and(|e| *e)
186 }
187
188 pub fn record_outbound(&self, message: Message) {
190 self.record(MessageDirection::Outbound, message);
191 }
192
193 pub fn record_inbound(&self, message: Message) {
195 self.record(MessageDirection::Inbound, message);
196 }
197
198 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 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 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 if let Ok(mut records) = self.records.write() {
227 records.push(record);
228
229 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 #[must_use]
239 pub fn records(&self) -> Vec<MessageRecord> {
240 self.records.read().map(|r| r.clone()).unwrap_or_default()
241 }
242
243 #[must_use]
245 pub fn len(&self) -> usize {
246 self.records.read().map_or(0, |r| r.len())
247 }
248
249 #[must_use]
251 pub fn is_empty(&self) -> bool {
252 self.len() == 0
253 }
254
255 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 #[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 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 #[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 #[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 #[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 #[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 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 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); }
446}