aether_telemetry/
telemetry_runtime.rs

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

aether_telemetry/
telemetry_runtime.rs
