aether_telemetry/
telemetry_runtime.rs

1use crate::error::{TelemetryInitError, TelemetryShutdownError};
2use crate::gen_ai_metrics::GenAiMetrics;
3use crate::genai_constants;
4use crate::otel_observer::OtelInstrumentation;
5use opentelemetry::KeyValue;
6use opentelemetry::metrics::MeterProvider as _;
7use opentelemetry::trace::TracerProvider as _;
8use opentelemetry_otlp::{MetricExporter, SpanExporter, WithExportConfig, WithHttpConfig};
9use opentelemetry_sdk::Resource;
10use opentelemetry_sdk::metrics::SdkMeterProvider;
11use opentelemetry_sdk::metrics::periodic_reader_with_async_runtime::PeriodicReader;
12use opentelemetry_sdk::runtime::Tokio;
13use opentelemetry_sdk::trace::span_processor_with_async_runtime::BatchSpanProcessor;
14use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
15use std::collections::HashMap;
16
17pub struct TelemetryRuntime {
18    tracer_provider: SdkTracerProvider,
19    meter_provider: SdkMeterProvider,
20    instrumentation: OtelInstrumentation,
21}
22
23/// Fully-resolved telemetry configuration. Every field is required: defaults
24/// and settings-layer merging live in `TelemetrySettings`, so a value here is
25/// always deliberate.
26pub struct TelemetryConfig {
27    /// Base URL for the OTLP/HTTP collector, required whenever a signal needs an exporter.
28    pub endpoint: Option<String>,
29    pub headers: HashMap<String, String>,
30    pub service_name: String,
31    pub service_version: String,
32    pub sample_ratio: f64,
33    pub capture_content: bool,
34    pub traces_enabled: bool,
35    pub metrics_enabled: bool,
36}
37
38impl TelemetryRuntime {
39    pub fn new(config: &TelemetryConfig) -> Result<Self, TelemetryInitError> {
40        if !(0.0..=1.0).contains(&config.sample_ratio) {
41            return Err(TelemetryInitError::InvalidSampleRatio(config.sample_ratio));
42        }
43
44        let http_client = build_http_client(&config.headers)?;
45        let resource = Resource::builder()
46            .with_service_name(config.service_name.clone())
47            .with_attribute(KeyValue::new(genai_constants::SERVICE_VERSION, config.service_version.clone()))
48            .build();
49        let scope = genai_constants::genai_instrumentation_scope(config.service_version.clone());
50        let tracer_provider = build_tracer_provider(config, resource.clone(), http_client.clone())?;
51        let meter_provider = build_meter_provider(config, resource, http_client)?;
52        let instrumentation = OtelInstrumentation {
53            tracer: tracer_provider.tracer_with_scope(scope.clone()),
54            metrics: GenAiMetrics::new(&meter_provider.meter_with_scope(scope)),
55            capture_content: config.capture_content,
56        };
57
58        Ok(Self { tracer_provider, meter_provider, instrumentation })
59    }
60
61    pub fn observer_factory(&self) -> aether_core::events::ObserverFactory {
62        let instrumentation = self.instrumentation.clone();
63        std::sync::Arc::new(move || Box::new(crate::OtelObserver::new(instrumentation.clone())))
64    }
65
66    /// Flushes and shuts down, logging any failure. For callers whose only
67    /// recovery is logging.
68    pub fn shutdown_or_log(&self) {
69        if let Err(error) = self.shutdown() {
70            tracing::error!("Failed to shutdown telemetry: {error}");
71        }
72    }
73
74    /// Flushes and shuts down every signal. Both providers are always attempted;
75    /// providers for disabled signals have no exporters.
76    pub fn shutdown(&self) -> Result<(), TelemetryShutdownError> {
77        let traces = self.tracer_provider.shutdown().map_err(TelemetryShutdownError::Trace);
78        let metrics = self.meter_provider.shutdown().map_err(TelemetryShutdownError::Metric);
79        traces.and(metrics)
80    }
81}
82
83impl TelemetryConfig {
84    fn signal_endpoint(&self, signal: &str) -> Result<String, TelemetryInitError> {
85        Ok(format!("{}/v1/{signal}", self.required_endpoint()?.trim_end_matches('/')))
86    }
87
88    fn required_endpoint(&self) -> Result<&str, TelemetryInitError> {
89        self.endpoint.as_deref().filter(|endpoint| !endpoint.is_empty()).ok_or(TelemetryInitError::MissingOtlpEndpoint)
90    }
91}
92
93fn build_tracer_provider(
94    config: &TelemetryConfig,
95    resource: Resource,
96    http_client: reqwest::Client,
97) -> Result<SdkTracerProvider, TelemetryInitError> {
98    let builder = SdkTracerProvider::builder().with_resource(resource);
99    if !config.traces_enabled {
100        return Ok(builder.with_sampler(Sampler::AlwaysOff).build());
101    }
102
103    let endpoint = config.signal_endpoint("traces")?;
104    let exporter = SpanExporter::builder()
105        .with_http()
106        .with_endpoint(endpoint)
107        .with_http_client(http_client)
108        .build()
109        .map_err(TelemetryInitError::TraceExporter)?;
110
111    Ok(builder
112        .with_sampler(Sampler::TraceIdRatioBased(config.sample_ratio))
113        .with_span_processor(BatchSpanProcessor::builder(exporter, Tokio).build())
114        .build())
115}
116
117fn build_meter_provider(
118    config: &TelemetryConfig,
119    resource: Resource,
120    http_client: reqwest::Client,
121) -> Result<SdkMeterProvider, TelemetryInitError> {
122    let builder = SdkMeterProvider::builder().with_resource(resource);
123    if !config.metrics_enabled {
124        return Ok(builder.build());
125    }
126
127    let endpoint = config.signal_endpoint("metrics")?;
128    let exporter = MetricExporter::builder()
129        .with_http()
130        .with_endpoint(endpoint)
131        .with_http_client(http_client)
132        .build()
133        .map_err(TelemetryInitError::MetricExporter)?;
134
135    Ok(builder.with_reader(PeriodicReader::builder(exporter, Tokio).build()).build())
136}
137
138fn build_http_client(headers: &HashMap<String, String>) -> Result<reqwest::Client, TelemetryInitError> {
139    let mut parsed = reqwest::header::HeaderMap::new();
140    for (key, value) in headers {
141        let name = reqwest::header::HeaderName::from_bytes(key.as_bytes())
142            .map_err(|_| TelemetryInitError::InvalidHeaderName(key.clone()))?;
143        let value = reqwest::header::HeaderValue::from_str(value)
144            .map_err(|_| TelemetryInitError::InvalidHeaderValue(key.clone()))?;
145        parsed.insert(name, value);
146    }
147    reqwest::Client::builder().default_headers(parsed).build().map_err(TelemetryInitError::HttpClient)
148}
aether_telemetry/telemetry_runtime.rs

aether_telemetry/
telemetry_runtime.rs
