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

aether_telemetry/
telemetry_runtime.rs
