Skip to main content

aptu_coder/
otel.rs

1// SPDX-FileCopyrightText: 2026 aptu-coder contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use opentelemetry::global;
5use opentelemetry_otlp::{Protocol, WithExportConfig};
6use opentelemetry_sdk::Resource;
7use opentelemetry_sdk::logs::SdkLoggerProvider;
8use opentelemetry_sdk::metrics::SdkMeterProvider;
9use opentelemetry_sdk::trace::SdkTracerProvider;
10use tracing::warn;
11
12/// Session and client metadata recorded as span attributes on every tool call.
13pub struct ClientMetadata {
14    pub session_id: Option<String>,
15    pub client_name: Option<String>,
16    pub client_version: Option<String>,
17}
18
19/// Extract W3C Trace Context from MCP request _meta field and set as parent span context.
20///
21/// Attempts to extract traceparent and tracestate from the request's _meta field.
22/// If successful, calls `set_parent` on the current tracing span so the OTel layer
23/// re-parents it to the caller's trace. This must be called after the `#[instrument]`
24/// span has been entered (i.e., inside the function body) for `set_parent` to take effect.
25/// If extraction fails or _meta is absent, silently proceeds with root context (no panic).
26pub fn extract_and_set_trace_context(
27    meta: Option<&rmcp::model::MetaObject>,
28    client_meta: ClientMetadata,
29) {
30    use tracing_opentelemetry::OpenTelemetrySpanExt as _;
31
32    let span = tracing::Span::current();
33
34    // Record session and client attributes
35    if let Some(sid) = client_meta.session_id {
36        span.record("mcp.session.id", &sid);
37    }
38    if let Some(cn) = client_meta.client_name {
39        span.record("client.name", &cn);
40    }
41    if let Some(cv) = client_meta.client_version {
42        span.record("client.version", &cv);
43    }
44
45    // Extract agent-session-id from _meta if present (opportunistic; silent no-op if absent)
46    if let Some(asi_str) = meta.and_then(|m| m.0.get("agent-session-id").and_then(|v| v.as_str())) {
47        span.record("mcp.client.session.id", asi_str);
48    }
49
50    let Some(meta) = meta else { return };
51
52    let mut propagation_map = std::collections::HashMap::new();
53
54    // Extract traceparent if present
55    if let Some(traceparent) = meta.0.get("traceparent")
56        && let Some(tp_str) = traceparent.as_str()
57    {
58        propagation_map.insert("traceparent".to_string(), tp_str.to_string());
59    }
60
61    // Extract tracestate if present
62    if let Some(tracestate) = meta.0.get("tracestate")
63        && let Some(ts_str) = tracestate.as_str()
64    {
65        propagation_map.insert("tracestate".to_string(), ts_str.to_string());
66    }
67
68    // Only attempt extraction if we have at least traceparent
69    if propagation_map.is_empty() {
70        return;
71    }
72
73    // Extract context via the globally registered propagator (TraceContextPropagator by default)
74    let parent_cx = opentelemetry::global::get_text_map_propagator(|propagator| {
75        propagator.extract(&ExtractMap(&propagation_map))
76    });
77
78    // Re-parent the current tracing span (already entered via #[instrument]) to the
79    // extracted OTel context. set_parent is a no-op if the OTel layer is not installed.
80    let _ = span.set_parent(parent_cx);
81}
82
83/// Helper struct for W3C Trace Context extraction from HashMap
84struct ExtractMap<'a>(&'a std::collections::HashMap<String, String>);
85
86impl<'a> opentelemetry::propagation::Extractor for ExtractMap<'a> {
87    fn get(&self, key: &str) -> Option<&str> {
88        self.0.get(key).map(|s| s.as_str())
89    }
90
91    fn keys(&self) -> Vec<&str> {
92        self.0.keys().map(|k| k.as_str()).collect()
93    }
94}
95
96/// Builds the standard service resource attached to all three signal providers.
97fn service_resource() -> Resource {
98    Resource::builder()
99        .with_attribute(opentelemetry::KeyValue::new("service.name", "aptu-coder"))
100        .with_attribute(opentelemetry::KeyValue::new(
101            "service.version",
102            env!("CARGO_PKG_VERSION"),
103        ))
104        .build()
105}
106
107/// Initializes OpenTelemetry with OTLP export if OTEL_EXPORTER_OTLP_ENDPOINT is set.
108///
109/// Returns `Some(provider)` if initialization succeeds, or `None` if:
110/// - The env var is unset (noop provider, zero overhead)
111/// - The exporter fails to build (logs warning, graceful failure)
112///
113/// The provider is registered globally via `opentelemetry::global::set_tracer_provider`.
114pub fn init_otel() -> Option<SdkTracerProvider> {
115    let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
116
117    // Build the OTLP exporter with HTTP proto transport
118    let exporter = match opentelemetry_otlp::SpanExporter::builder()
119        .with_http()
120        .with_protocol(Protocol::HttpBinary)
121        .with_endpoint(&endpoint)
122        .build()
123    {
124        Ok(exp) => exp,
125        Err(e) => {
126            warn!("Failed to build OTLP exporter: {}", e);
127            return None;
128        }
129    };
130
131    // Build provider with batch exporter for async export
132    let provider = SdkTracerProvider::builder()
133        .with_batch_exporter(exporter)
134        .with_resource(service_resource())
135        .build();
136
137    // Register globally
138    global::set_tracer_provider(provider.clone());
139
140    Some(provider)
141}
142
143/// Initializes OpenTelemetry log appender if OTEL_EXPORTER_OTLP_ENDPOINT is set.
144///
145/// Returns `Some(provider)` if initialization succeeds, or `None` if:
146/// - The env var is unset (noop, zero overhead)
147/// - The exporter fails to build (logs warning, graceful failure)
148///
149/// The provider is returned for use with OpenTelemetryTracingBridge layer.
150pub fn init_log_appender() -> Option<SdkLoggerProvider> {
151    let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
152
153    // Build the OTLP log exporter with HTTP proto transport
154    let exporter = match opentelemetry_otlp::LogExporter::builder()
155        .with_http()
156        .with_protocol(Protocol::HttpBinary)
157        .with_endpoint(&endpoint)
158        .build()
159    {
160        Ok(exp) => exp,
161        Err(e) => {
162            warn!("Failed to build OTLP log exporter: {}", e);
163            return None;
164        }
165    };
166
167    // Build provider with batch processor for async export
168    let provider = SdkLoggerProvider::builder()
169        .with_batch_exporter(exporter)
170        .with_resource(service_resource())
171        .build();
172
173    Some(provider)
174}
175
176/// Initializes OpenTelemetry metrics SDK if OTEL_EXPORTER_OTLP_ENDPOINT is set.
177///
178/// Returns `Some(provider)` if initialization succeeds, or `None` if:
179/// - The env var is unset (noop, zero overhead)
180/// - The exporter fails to build (logs warning, graceful failure)
181///
182/// The provider is registered globally via `opentelemetry::global::set_meter_provider`.
183pub fn init_meter() -> Option<SdkMeterProvider> {
184    let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
185
186    // Build the OTLP metrics exporter with HTTP proto transport
187    let exporter = match opentelemetry_otlp::MetricExporter::builder()
188        .with_http()
189        .with_protocol(Protocol::HttpBinary)
190        .with_endpoint(&endpoint)
191        .build()
192    {
193        Ok(exp) => exp,
194        Err(e) => {
195            warn!("Failed to build OTLP metrics exporter: {}", e);
196            return None;
197        }
198    };
199
200    // Build provider with periodic reader for async export
201    let provider = SdkMeterProvider::builder()
202        .with_reader(opentelemetry_sdk::metrics::PeriodicReader::builder(exporter).build())
203        .with_resource(service_resource())
204        .build();
205
206    // Register globally
207    global::set_meter_provider(provider.clone());
208
209    Some(provider)
210}