1use 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
12pub struct ClientMetadata {
14 pub session_id: Option<String>,
15 pub client_name: Option<String>,
16 pub client_version: Option<String>,
17}
18
19pub 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 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 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 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 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 if propagation_map.is_empty() {
70 return;
71 }
72
73 let parent_cx = opentelemetry::global::get_text_map_propagator(|propagator| {
75 propagator.extract(&ExtractMap(&propagation_map))
76 });
77
78 let _ = span.set_parent(parent_cx);
81}
82
83struct 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
96fn 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
107pub fn init_otel() -> Option<SdkTracerProvider> {
115 let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
116
117 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 let provider = SdkTracerProvider::builder()
133 .with_batch_exporter(exporter)
134 .with_resource(service_resource())
135 .build();
136
137 global::set_tracer_provider(provider.clone());
139
140 Some(provider)
141}
142
143pub fn init_log_appender() -> Option<SdkLoggerProvider> {
151 let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
152
153 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 let provider = SdkLoggerProvider::builder()
169 .with_batch_exporter(exporter)
170 .with_resource(service_resource())
171 .build();
172
173 Some(provider)
174}
175
176pub fn init_meter() -> Option<SdkMeterProvider> {
184 let endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok()?;
185
186 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 let provider = SdkMeterProvider::builder()
202 .with_reader(opentelemetry_sdk::metrics::PeriodicReader::builder(exporter).build())
203 .with_resource(service_resource())
204 .build();
205
206 global::set_meter_provider(provider.clone());
208
209 Some(provider)
210}