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, 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::{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 root_parent = extract_root_parent(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 = build_tracer_provider(config, resource.clone(), http_client.clone())?;
63        let meter_provider = build_meter_provider(config, resource, http_client)?;
64        let instrumentation = OtelInstrumentation {
65            tracer: tracer_provider.tracer_with_scope(scope.clone()),
66            metrics: GenAiMetrics::new(&meter_provider.meter_with_scope(scope)),
67            capture_content: config.capture_content,
68            root_parent,
69        };
70
71        Ok(Self { tracer_provider, meter_provider, instrumentation })
72    }
73
74    pub fn observer_factory(&self) -> aether_core::events::ObserverFactory {
75        let instrumentation = self.instrumentation.clone();
76        std::sync::Arc::new(move || Box::new(crate::OtelObserver::new(instrumentation.clone())))
77    }
78
79    /// Flushes and shuts down, logging any failure. For callers whose only
80    /// recovery is logging.
81    pub fn shutdown_or_log(&self) {
82        if let Err(error) = self.shutdown() {
83            tracing::error!("Failed to shutdown telemetry: {error}");
84        }
85    }
86
87    /// Flushes and shuts down every signal. Both providers are always attempted;
88    /// providers for disabled signals have no exporters.
89    pub fn shutdown(&self) -> Result<(), TelemetryShutdownError> {
90        let traces = self.tracer_provider.shutdown().map_err(TelemetryShutdownError::Trace);
91        let metrics = self.meter_provider.shutdown().map_err(TelemetryShutdownError::Metric);
92        traces.and(metrics)
93    }
94}
95
96impl TelemetryConfig {
97    fn signal_endpoint(&self, signal: &str) -> Result<String, TelemetryInitError> {
98        let signal_specific_endpoint = match signal {
99            "traces" => self.traces_endpoint.as_deref(),
100            "metrics" => self.metrics_endpoint.as_deref(),
101            _ => None,
102        };
103        if let Some(endpoint) = signal_specific_endpoint.filter(|endpoint| !endpoint.is_empty()) {
104            return Ok(endpoint.to_string());
105        }
106
107        Ok(format!("{}/v1/{signal}", self.required_endpoint()?.trim_end_matches('/')))
108    }
109
110    fn required_endpoint(&self) -> Result<&str, TelemetryInitError> {
111        self.endpoint.as_deref().filter(|endpoint| !endpoint.is_empty()).ok_or(TelemetryInitError::MissingOtlpEndpoint)
112    }
113}
114
115fn extract_root_parent(trace_context: Option<&AgentTraceContext>) -> Result<Option<Context>, TelemetryInitError> {
116    let Some(trace_context) = trace_context else {
117        return Ok(None);
118    };
119
120    if let Some(tracestate) = &trace_context.tracestate {
121        TraceState::from_str(tracestate).map_err(|_| TelemetryInitError::InvalidTraceContext("tracestate"))?;
122    }
123
124    let context = TraceContextPropagator::new().extract(trace_context);
125    if !context.span().span_context().is_valid() {
126        return Err(TelemetryInitError::InvalidTraceContext("traceparent"));
127    }
128
129    Ok(Some(context))
130}
131
132fn build_tracer_provider(
133    config: &TelemetryConfig,
134    resource: Resource,
135    http_client: reqwest::Client,
136) -> Result<SdkTracerProvider, TelemetryInitError> {
137    let builder = SdkTracerProvider::builder().with_resource(resource);
138    if !config.traces_enabled {
139        return Ok(builder.with_sampler(Sampler::AlwaysOff).build());
140    }
141
142    let endpoint = config.signal_endpoint("traces")?;
143    let exporter = SpanExporter::builder()
144        .with_http()
145        .with_endpoint(endpoint)
146        .with_http_client(http_client)
147        .build()
148        .map_err(TelemetryInitError::TraceExporter)?;
149
150    Ok(builder
151        .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(config.sample_ratio))))
152        .with_span_processor(BatchSpanProcessor::builder(exporter, Tokio).build())
153        .build())
154}
155
156fn build_meter_provider(
157    config: &TelemetryConfig,
158    resource: Resource,
159    http_client: reqwest::Client,
160) -> Result<SdkMeterProvider, TelemetryInitError> {
161    let builder = SdkMeterProvider::builder().with_resource(resource);
162    if !config.metrics_enabled {
163        return Ok(builder.build());
164    }
165
166    let endpoint = config.signal_endpoint("metrics")?;
167    let exporter = MetricExporter::builder()
168        .with_http()
169        .with_endpoint(endpoint)
170        .with_http_client(http_client)
171        .build()
172        .map_err(TelemetryInitError::MetricExporter)?;
173
174    Ok(builder.with_reader(PeriodicReader::builder(exporter, Tokio).build()).build())
175}
176
177fn build_http_client(headers: &HashMap<String, String>) -> Result<reqwest::Client, TelemetryInitError> {
178    let mut parsed = reqwest::header::HeaderMap::new();
179    for (key, value) in headers {
180        let name = reqwest::header::HeaderName::from_bytes(key.as_bytes())
181            .map_err(|_| TelemetryInitError::InvalidHeaderName(key.clone()))?;
182        let value = reqwest::header::HeaderValue::from_str(value)
183            .map_err(|_| TelemetryInitError::InvalidHeaderValue(key.clone()))?;
184        parsed.insert(name, value);
185    }
186    reqwest::Client::builder().default_headers(parsed).build().map_err(TelemetryInitError::HttpClient)
187}
aether_telemetry/telemetry_runtime.rs

aether_telemetry/
telemetry_runtime.rs
