Skip to main content

prism_mcp_rs/
telemetry.rs

1//! OpenTelemetry setup for exporting the request spans emitted by the SDK.
2//!
3//! Enable the `otel` feature and call [`init_otlp_tracing`] once at process
4//! startup. The server request path is instrumented even when this feature is
5//! disabled, so applications may also install their own tracing subscriber.
6
7use crate::core::error::{McpError, McpResult};
8use opentelemetry::global;
9use opentelemetry_otlp::WithExportConfig;
10use opentelemetry_sdk::propagation::TraceContextPropagator;
11use opentelemetry_sdk::trace::SdkTracerProvider;
12use opentelemetry_sdk::Resource;
13use tracing_subscriber::layer::SubscriberExt;
14use tracing_subscriber::util::SubscriberInitExt;
15use tracing_subscriber::EnvFilter;
16
17/// OTLP trace export configuration.
18#[derive(Debug, Clone)]
19pub struct OtlpTracingConfig {
20    pub service_name: String,
21    pub endpoint: String,
22    pub filter: String,
23}
24
25impl OtlpTracingConfig {
26    pub fn new(service_name: impl Into<String>, endpoint: impl Into<String>) -> Self {
27        Self {
28            service_name: service_name.into(),
29            endpoint: endpoint.into(),
30            filter: "info".to_string(),
31        }
32    }
33
34    pub fn with_filter(mut self, filter: impl Into<String>) -> Self {
35        self.filter = filter.into();
36        self
37    }
38}
39
40/// Flushes batched spans when dropped or explicitly shut down.
41pub struct TelemetryGuard {
42    provider: Option<SdkTracerProvider>,
43}
44
45impl TelemetryGuard {
46    pub fn shutdown(mut self) -> McpResult<()> {
47        self.provider
48            .take()
49            .expect("telemetry provider is present until shutdown")
50            .shutdown()
51            .map_err(|error| McpError::Internal(format!("OpenTelemetry shutdown failed: {error}")))
52    }
53}
54
55impl Drop for TelemetryGuard {
56    fn drop(&mut self) {
57        if let Some(provider) = self.provider.take() {
58            if let Err(error) = provider.shutdown() {
59                tracing::warn!(%error, "OpenTelemetry provider did not shut down cleanly");
60            }
61        }
62    }
63}
64
65/// Install an OTLP/gRPC exporter and a W3C Trace Context propagator.
66///
67/// This function must be called from within a Tokio runtime and before another
68/// global tracing subscriber has been installed.
69pub fn init_otlp_tracing(config: OtlpTracingConfig) -> McpResult<TelemetryGuard> {
70    let exporter = opentelemetry_otlp::SpanExporter::builder()
71        .with_tonic()
72        .with_endpoint(config.endpoint)
73        .build()
74        .map_err(|error| McpError::Internal(format!("failed to build OTLP exporter: {error}")))?;
75
76    let resource = Resource::builder()
77        .with_service_name(config.service_name)
78        .build();
79    let provider = SdkTracerProvider::builder()
80        .with_batch_exporter(exporter)
81        .with_resource(resource)
82        .build();
83
84    global::set_text_map_propagator(TraceContextPropagator::new());
85    global::set_tracer_provider(provider.clone());
86    let tracer = global::tracer("prism-mcp-rs");
87    let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer);
88    let filter = EnvFilter::try_new(config.filter)
89        .map_err(|error| McpError::Validation(format!("invalid tracing filter: {error}")))?;
90
91    tracing_subscriber::registry()
92        .with(filter)
93        .with(tracing_subscriber::fmt::layer())
94        .with(telemetry_layer)
95        .try_init()
96        .map_err(|error| {
97            McpError::Internal(format!("failed to install tracing subscriber: {error}"))
98        })?;
99
100    Ok(TelemetryGuard {
101        provider: Some(provider),
102    })
103}