aether_telemetry/
telemetry_runtime.rs

1use crate::AgentTraceContext;
2use crate::error::{TelemetryInitError, TelemetryShutdownError};
3use crate::gen_ai_metrics::GenAiMetrics;
4use crate::genai_constants;
5use crate::otel_observer::OtelInstrumentation;
6use crate::span_guard::{ErrorKind, SpanGuard};
7use crate::trace_context::{extract_trace_context, inject_trace_context};
8use aether_core::events::{
9    AgentObserver, DynObserverFactory, McpRequestInstrumentation, ObserverFactory, TraceContext,
10};
11use opentelemetry::metrics::MeterProvider as _;
12use opentelemetry::trace::{SpanBuilder, SpanKind, TraceId, TraceState, TracerProvider as _};
13use opentelemetry::{Context, KeyValue};
14use opentelemetry_otlp::{MetricExporter, SpanExporter, WithExportConfig, WithHttpConfig};
15use opentelemetry_sdk::Resource;
16use opentelemetry_sdk::metrics::SdkMeterProvider;
17use opentelemetry_sdk::metrics::periodic_reader_with_async_runtime::PeriodicReader;
18use opentelemetry_sdk::runtime::Tokio;
19use opentelemetry_sdk::trace::span_processor_with_async_runtime::BatchSpanProcessor;
20use opentelemetry_sdk::trace::{IdGenerator, RandomIdGenerator, Sampler, SdkTracerProvider};
21use std::collections::HashMap;
22use std::str::FromStr;
23
24pub struct TelemetryRuntime {
25    tracer_provider: SdkTracerProvider,
26    meter_provider: SdkMeterProvider,
27    instrumentation: OtelInstrumentation,
28}
29
30struct OtelObserverFactory {
31    instrumentation: OtelInstrumentation,
32}
33
34struct OtelMcpRequestInstrumentation {
35    span: SpanGuard,
36}
37
38/// Fully-resolved telemetry configuration. Every field is required: defaults
39/// and settings-layer merging live in `TelemetrySettings`, so a value here is
40/// always deliberate.
41pub struct TelemetryConfig {
42    /// Base URL for the OTLP/HTTP collector, required whenever a signal needs an exporter.
43    pub endpoint: Option<String>,
44    /// Exact OTLP/HTTP trace export URL. When set, this takes precedence over the
45    /// trace URL derived from `endpoint`.
46    pub traces_endpoint: Option<String>,
47    /// Exact OTLP/HTTP metric export URL. When set, this takes precedence over the
48    /// metric URL derived from `endpoint`.
49    pub metrics_endpoint: Option<String>,
50    pub headers: HashMap<String, String>,
51    pub service_name: String,
52    pub service_version: String,
53    pub sample_ratio: f64,
54    pub capture_content: bool,
55    pub trace_context: Option<AgentTraceContext>,
56    pub traces_enabled: bool,
57    pub metrics_enabled: bool,
58}
59
60impl TelemetryRuntime {
61    pub fn new(config: &TelemetryConfig) -> Result<Self, TelemetryInitError> {
62        if !(0.0..=1.0).contains(&config.sample_ratio) {
63            return Err(TelemetryInitError::InvalidSampleRatio(config.sample_ratio));
64        }
65
66        let trace_root = resolve_trace_root(config.trace_context.as_ref())?;
67        let http_client = build_http_client(&config.headers)?;
68        let resource = Resource::builder()
69            .with_service_name(config.service_name.clone())
70            .with_attribute(KeyValue::new(genai_constants::SERVICE_VERSION, config.service_version.clone()))
71            .build();
72        let scope = genai_constants::genai_instrumentation_scope(config.service_version.clone());
73        let tracer_provider =
74            build_tracer_provider(config, resource.clone(), http_client.clone(), trace_root.trace_id)?;
75        let meter_provider = build_meter_provider(config, resource, http_client)?;
76        let instrumentation = OtelInstrumentation {
77            tracer: tracer_provider.tracer_with_scope(scope.clone()),
78            metrics: GenAiMetrics::new(&meter_provider.meter_with_scope(scope)),
79            capture_content: config.capture_content,
80            root_parent: trace_root.parent,
81        };
82
83        Ok(Self { tracer_provider, meter_provider, instrumentation })
84    }
85
86    pub fn observer_factory(&self) -> DynObserverFactory {
87        std::sync::Arc::new(OtelObserverFactory { instrumentation: self.instrumentation.clone() })
88    }
89
90    /// Flushes and shuts down, logging any failure. For callers whose only
91    /// recovery is logging.
92    pub fn shutdown_or_log(&self) {
93        if let Err(error) = self.shutdown() {
94            tracing::error!("Failed to shutdown telemetry: {error}");
95        }
96    }
97
98    /// Flushes and shuts down every signal. Both providers are always attempted;
99    /// providers for disabled signals have no exporters.
100    pub fn shutdown(&self) -> Result<(), TelemetryShutdownError> {
101        let traces = self.tracer_provider.shutdown().map_err(TelemetryShutdownError::Trace);
102        let metrics = self.meter_provider.shutdown().map_err(TelemetryShutdownError::Metric);
103        traces.and(metrics)
104    }
105}
106
107impl ObserverFactory for OtelObserverFactory {
108    fn agent(&self, parent: Option<&TraceContext>) -> Box<dyn AgentObserver> {
109        let mut instrumentation = self.instrumentation.clone();
110        if let Some(parent) = parent.and_then(extract_trace_context) {
111            instrumentation.root_parent = Some(parent);
112        }
113        Box::new(crate::OtelObserver::new(instrumentation))
114    }
115
116    fn tool_call_request(&self, tool_name: &str, parent: Option<&TraceContext>) -> Box<dyn McpRequestInstrumentation> {
117        let remote_parent = parent.and_then(extract_trace_context);
118        let attributes = vec![
119            KeyValue::new(genai_constants::GEN_AI_OPERATION_NAME, "execute_tool"),
120            KeyValue::new(genai_constants::GEN_AI_TOOL_NAME, tool_name.to_string()),
121            KeyValue::new(genai_constants::MCP_METHOD_NAME, TOOLS_CALL_METHOD),
122            KeyValue::new(genai_constants::MCP_TOOL_NAME, tool_name.to_string()),
123        ];
124        let builder = SpanBuilder::from_name(format!("{TOOLS_CALL_METHOD} {tool_name}"))
125            .with_kind(SpanKind::Server)
126            .with_attributes(attributes);
127        let context = self
128            .instrumentation
129            .start_span(builder, remote_parent.as_ref().or(self.instrumentation.root_parent.as_ref()));
130        Box::new(OtelMcpRequestInstrumentation { span: SpanGuard::new(context, "MCP request cancelled") })
131    }
132}
133
134impl McpRequestInstrumentation for OtelMcpRequestInstrumentation {
135    fn trace_context(&self) -> Option<TraceContext> {
136        inject_trace_context(self.span.context())
137    }
138
139    fn finish(mut self: Box<Self>, error: Option<&str>) {
140        match error {
141            Some(error) => self.span.end_error(Some(ErrorKind::McpError), error),
142            None => self.span.end_ok(),
143        }
144    }
145}
146
147impl TelemetryConfig {
148    fn signal_endpoint(&self, signal: &str) -> Result<String, TelemetryInitError> {
149        let signal_specific_endpoint = match signal {
150            "traces" => self.traces_endpoint.as_deref(),
151            "metrics" => self.metrics_endpoint.as_deref(),
152            _ => None,
153        };
154        if let Some(endpoint) = signal_specific_endpoint.filter(|endpoint| !endpoint.is_empty()) {
155            return Ok(endpoint.to_string());
156        }
157
158        Ok(format!("{}/v1/{signal}", self.required_endpoint()?.trim_end_matches('/')))
159    }
160
161    fn required_endpoint(&self) -> Result<&str, TelemetryInitError> {
162        self.endpoint.as_deref().filter(|endpoint| !endpoint.is_empty()).ok_or(TelemetryInitError::MissingOtlpEndpoint)
163    }
164}
165
166const TOOLS_CALL_METHOD: &str = "tools/call";
167
168#[derive(Default)]
169struct TraceRoot {
170    parent: Option<Context>,
171    trace_id: Option<TraceId>,
172}
173
174#[derive(Debug)]
175struct TraceIdGenerator {
176    trace_id: TraceId,
177    random: RandomIdGenerator,
178}
179
180impl IdGenerator for TraceIdGenerator {
181    fn new_trace_id(&self) -> TraceId {
182        self.trace_id
183    }
184
185    fn new_span_id(&self) -> opentelemetry::trace::SpanId {
186        self.random.new_span_id()
187    }
188}
189
190fn resolve_trace_root(trace_context: Option<&AgentTraceContext>) -> Result<TraceRoot, TelemetryInitError> {
191    let Some(trace_context) = trace_context else {
192        return Ok(TraceRoot::default());
193    };
194
195    match trace_context {
196        AgentTraceContext::Parent(trace_context) => {
197            if let Some(tracestate) = &trace_context.tracestate {
198                TraceState::from_str(tracestate).map_err(|_| TelemetryInitError::InvalidTraceContext("tracestate"))?;
199            }
200
201            let parent =
202                extract_trace_context(trace_context).ok_or(TelemetryInitError::InvalidTraceContext("traceparent"))?;
203
204            Ok(TraceRoot { parent: Some(parent), trace_id: None })
205        }
206        AgentTraceContext::Root { trace_id } => {
207            let trace_id = parse_trace_id(trace_id)?;
208            Ok(TraceRoot { parent: None, trace_id: Some(trace_id) })
209        }
210    }
211}
212
213fn parse_trace_id(value: &str) -> Result<TraceId, TelemetryInitError> {
214    let valid_hex =
215        value.len() == 32 && value.bytes().all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte));
216    if !valid_hex {
217        return Err(TelemetryInitError::InvalidTraceContext("traceId"));
218    }
219
220    let trace_id = TraceId::from_hex(value).map_err(|_| TelemetryInitError::InvalidTraceContext("traceId"))?;
221    if trace_id == TraceId::INVALID {
222        return Err(TelemetryInitError::InvalidTraceContext("traceId"));
223    }
224
225    Ok(trace_id)
226}
227
228fn build_tracer_provider(
229    config: &TelemetryConfig,
230    resource: Resource,
231    http_client: reqwest::Client,
232    trace_id: Option<TraceId>,
233) -> Result<SdkTracerProvider, TelemetryInitError> {
234    let builder = SdkTracerProvider::builder().with_resource(resource);
235    let builder = match trace_id {
236        Some(trace_id) => {
237            builder.with_id_generator(TraceIdGenerator { trace_id, random: RandomIdGenerator::default() })
238        }
239        None => builder,
240    };
241    if !config.traces_enabled {
242        return Ok(builder.with_sampler(Sampler::AlwaysOff).build());
243    }
244
245    let endpoint = config.signal_endpoint("traces")?;
246    let exporter = SpanExporter::builder()
247        .with_http()
248        .with_endpoint(endpoint)
249        .with_http_client(http_client)
250        .build()
251        .map_err(TelemetryInitError::TraceExporter)?;
252
253    Ok(builder
254        .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(config.sample_ratio))))
255        .with_span_processor(BatchSpanProcessor::builder(exporter, Tokio).build())
256        .build())
257}
258
259fn build_meter_provider(
260    config: &TelemetryConfig,
261    resource: Resource,
262    http_client: reqwest::Client,
263) -> Result<SdkMeterProvider, TelemetryInitError> {
264    let builder = SdkMeterProvider::builder().with_resource(resource);
265    if !config.metrics_enabled {
266        return Ok(builder.build());
267    }
268
269    let endpoint = config.signal_endpoint("metrics")?;
270    let exporter = MetricExporter::builder()
271        .with_http()
272        .with_endpoint(endpoint)
273        .with_http_client(http_client)
274        .build()
275        .map_err(TelemetryInitError::MetricExporter)?;
276
277    Ok(builder.with_reader(PeriodicReader::builder(exporter, Tokio).build()).build())
278}
279
280fn build_http_client(headers: &HashMap<String, String>) -> Result<reqwest::Client, TelemetryInitError> {
281    let mut parsed = reqwest::header::HeaderMap::new();
282    for (key, value) in headers {
283        let name = reqwest::header::HeaderName::from_bytes(key.as_bytes())
284            .map_err(|_| TelemetryInitError::InvalidHeaderName(key.clone()))?;
285        let value = reqwest::header::HeaderValue::from_str(value)
286            .map_err(|_| TelemetryInitError::InvalidHeaderValue(key.clone()))?;
287        parsed.insert(name, value);
288    }
289    reqwest::Client::builder().default_headers(parsed).build().map_err(TelemetryInitError::HttpClient)
290}
aether_telemetry/telemetry_runtime.rs

aether_telemetry/
telemetry_runtime.rs
