Skip to main content

code_analyze_mcp/
logging.rs

1use rmcp::model::LoggingLevel;
2use serde_json::{Map, Value};
3use std::sync::{Arc, Mutex};
4use tokio::sync::mpsc;
5use tracing::span::Attributes;
6use tracing::subscriber::Interest;
7use tracing::{Level, Subscriber};
8use tracing_subscriber::Layer;
9use tracing_subscriber::filter::LevelFilter;
10use tracing_subscriber::layer::Context;
11
12/// Maps tracing::Level to MCP LoggingLevel.
13pub fn level_to_mcp(level: &Level) -> LoggingLevel {
14    match *level {
15        Level::TRACE | Level::DEBUG => LoggingLevel::Debug,
16        Level::INFO => LoggingLevel::Info,
17        Level::WARN => LoggingLevel::Warning,
18        Level::ERROR => LoggingLevel::Error,
19    }
20}
21
22/// Lightweight event sent from McpLoggingLayer to consumer task via unbounded channel.
23#[derive(Clone, Debug)]
24pub struct LogEvent {
25    pub level: LoggingLevel,
26    pub logger: String,
27    pub data: Value,
28}
29
30/// Custom tracing Layer that bridges tracing events to MCP client via unbounded channel.
31/// Sends lightweight LogEvent to channel; consumer task in on_initialized drains with recv_many.
32pub struct McpLoggingLayer {
33    event_tx: mpsc::UnboundedSender<LogEvent>,
34    log_level_filter: Arc<Mutex<LevelFilter>>,
35}
36
37impl McpLoggingLayer {
38    pub fn new(
39        event_tx: mpsc::UnboundedSender<LogEvent>,
40        log_level_filter: Arc<Mutex<LevelFilter>>,
41    ) -> Self {
42        Self {
43            event_tx,
44            log_level_filter,
45        }
46    }
47}
48
49impl<S> Layer<S> for McpLoggingLayer
50where
51    S: Subscriber,
52{
53    fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
54        let metadata = event.metadata();
55        let level = *metadata.level();
56        let target = metadata.target();
57
58        // Check if event level passes the current filter before processing
59        let filter_level = self.log_level_filter.lock().unwrap();
60        if level > *filter_level {
61            return;
62        }
63        drop(filter_level);
64
65        // Extract fields from the event using a visitor that collects into a Map.
66        let mut fields = Map::new();
67        let mut visitor = MessageVisitor(&mut fields);
68        event.record(&mut visitor);
69
70        let mcp_level = level_to_mcp(&level);
71        let logger = target.to_string();
72        let data = Value::Object(fields);
73
74        // Send LogEvent to channel without blocking on_event.
75        let log_event = LogEvent {
76            level: mcp_level,
77            logger,
78            data,
79        };
80
81        // Ignore send error if receiver is dropped (channel closed).
82        let _ = self.event_tx.send(log_event);
83    }
84
85    fn register_callsite(&self, metadata: &'static tracing::Metadata<'static>) -> Interest {
86        let filter_level = self.log_level_filter.lock().unwrap();
87        if *metadata.level() <= *filter_level {
88            Interest::always()
89        } else {
90            Interest::never()
91        }
92    }
93
94    fn enabled(&self, metadata: &tracing::Metadata<'_>, _ctx: Context<'_, S>) -> bool {
95        let filter_level = self.log_level_filter.lock().unwrap();
96        *metadata.level() <= *filter_level
97    }
98
99    fn on_new_span(&self, _attrs: &Attributes<'_>, _id: &tracing::span::Id, _ctx: Context<'_, S>) {}
100}
101
102/// Visitor to extract fields from tracing event into a JSON Map.
103struct MessageVisitor<'a>(&'a mut Map<String, Value>);
104
105impl<'a> tracing::field::Visit for MessageVisitor<'a> {
106    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
107        self.0.insert(
108            field.name().to_string(),
109            Value::String(format!("{:?}", value)),
110        );
111    }
112
113    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
114        self.0
115            .insert(field.name().to_string(), Value::String(value.to_string()));
116    }
117
118    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
119        self.0
120            .insert(field.name().to_string(), Value::Number(value.into()));
121    }
122
123    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
124        self.0
125            .insert(field.name().to_string(), Value::Number(value.into()));
126    }
127
128    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
129        self.0.insert(field.name().to_string(), Value::Bool(value));
130    }
131}