Skip to main content

adk_telemetry/
init.rs

1//! Telemetry initialization and configuration
2
3use std::sync::{Arc, Once};
4use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
5
6use crate::span_exporter::{AdkSpanExporter, AdkSpanLayer};
7
8static INIT: Once = Once::new();
9
10/// Error returned by telemetry initialization functions.
11#[derive(Debug, thiserror::Error)]
12pub enum TelemetryError {
13    /// Failed to build the tracing/OTLP pipeline.
14    #[error("telemetry init failed: {0}")]
15    Init(String),
16}
17
18/// Initialize basic telemetry with console logging.
19///
20/// # Arguments
21/// * `service_name` - Name of the service for trace identification
22///
23/// # Example
24/// ```
25/// use adk_telemetry::init_telemetry;
26/// init_telemetry("my-agent-service").expect("Failed to initialize telemetry");
27/// ```
28pub fn init_telemetry(service_name: &str) -> Result<(), TelemetryError> {
29    INIT.call_once(|| {
30        let filter = EnvFilter::try_from_default_env()
31            .or_else(|_| EnvFilter::try_new("info"))
32            .unwrap_or_else(|_| EnvFilter::new("info"));
33
34        tracing_subscriber::registry()
35            .with(filter)
36            .with(
37                tracing_subscriber::fmt::layer()
38                    .with_target(true)
39                    .with_thread_ids(true)
40                    .with_line_number(true),
41            )
42            .init();
43
44        tracing::info!(service.name = service_name, "telemetry initialized");
45    });
46
47    Ok(())
48}
49
50/// Initialize telemetry with OpenTelemetry OTLP export.
51///
52/// Enables distributed tracing by exporting spans to an OTLP collector.
53///
54/// # Arguments
55/// * `service_name` - Name of the service for trace identification
56/// * `endpoint` - OTLP collector endpoint (e.g., "http://localhost:4317")
57///
58/// # Example
59/// ```no_run
60/// use adk_telemetry::init_with_otlp;
61/// init_with_otlp("my-agent", "http://localhost:4317")
62///     .expect("Failed to initialize telemetry");
63/// ```
64#[cfg(feature = "otlp")]
65pub fn init_with_otlp(service_name: &str, endpoint: &str) -> Result<(), TelemetryError> {
66    use opentelemetry::trace::TracerProvider;
67    use opentelemetry_otlp::WithExportConfig;
68    use tracing_opentelemetry::OpenTelemetryLayer;
69
70    let endpoint = endpoint.to_string();
71    let service_name = service_name.to_string();
72
73    let init_error: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
74
75    INIT.call_once(|| {
76        let resource = opentelemetry_sdk::Resource::builder_empty()
77            .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.clone())])
78            .build();
79
80        // Build OTLP span exporter
81        let span_exporter = match opentelemetry_otlp::SpanExporter::builder()
82            .with_tonic()
83            .with_endpoint(&endpoint)
84            .build()
85        {
86            Ok(e) => e,
87            Err(e) => {
88                *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
89                    Some(format!("failed to build OTLP span exporter: {e}"));
90                return;
91            }
92        };
93
94        // Build tracer provider with batch exporter
95        let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
96            .with_batch_exporter(span_exporter)
97            .with_resource(resource.clone())
98            .build();
99
100        let tracer = tracer_provider.tracer("adk-telemetry");
101        opentelemetry::global::set_tracer_provider(tracer_provider);
102
103        // Initialize metrics
104        let metric_exporter = match opentelemetry_otlp::MetricExporter::builder()
105            .with_tonic()
106            .with_endpoint(&endpoint)
107            .build()
108        {
109            Ok(e) => e,
110            Err(e) => {
111                *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
112                    Some(format!("failed to build OTLP metric exporter: {e}"));
113                return;
114            }
115        };
116
117        let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
118            .with_periodic_exporter(metric_exporter)
119            .with_resource(resource)
120            .build();
121
122        opentelemetry::global::set_meter_provider(meter_provider);
123
124        let telemetry_layer = OpenTelemetryLayer::new(tracer);
125
126        let filter = EnvFilter::try_from_default_env()
127            .or_else(|_| EnvFilter::try_new("info"))
128            .unwrap_or_else(|_| EnvFilter::new("info"));
129
130        tracing_subscriber::registry()
131            .with(filter)
132            .with(
133                tracing_subscriber::fmt::layer()
134                    .with_target(true)
135                    .with_thread_ids(true)
136                    .with_line_number(true),
137            )
138            .with(telemetry_layer)
139            .init();
140
141        tracing::info!(
142            service.name = service_name,
143            otlp.endpoint = %endpoint,
144            "telemetry initialized with OpenTelemetry"
145        );
146    });
147
148    if let Some(err) = init_error.lock().unwrap_or_else(|p| p.into_inner()).take() {
149        return Err(TelemetryError::Init(err));
150    }
151
152    Ok(())
153}
154
155/// Build an OTLP tracing layer without initializing a global subscriber.
156///
157/// Returns a boxed [`tracing_subscriber::Layer`] that can be composed with any
158/// subscriber via `.with()`. Also configures the global OpenTelemetry tracer
159/// and meter providers.
160///
161/// The layer is returned as `Box<dyn Layer<S>>` rather than `impl Layer` so it
162/// can be stored, composed across crate boundaries, and used in `Layered<...>`
163/// chains without running into opaque-type limitations.
164///
165/// Unlike [`init_with_otlp`], this function does **not** call `.init()` on a
166/// subscriber and does **not** use the `INIT` [`Once`] guard. The caller is
167/// responsible for composing the returned layer into their own subscriber stack.
168///
169/// # Arguments
170/// * `service_name` - Name of the service for trace identification
171/// * `endpoint` - OTLP collector endpoint (e.g., `"http://localhost:4317"`)
172///
173/// # Errors
174/// Returns [`TelemetryError::Init`] if the OTLP span or metric exporter fails to build.
175///
176/// # Example
177/// ```no_run
178/// use adk_telemetry::build_otlp_layer;
179/// use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
180///
181/// let otlp_layer = build_otlp_layer("my-agent", "http://localhost:4317")
182///     .expect("Failed to build OTLP layer");
183///
184/// tracing_subscriber::registry()
185///     .with(otlp_layer)
186///     .with(tracing_subscriber::fmt::layer())
187///     .init();
188/// ```
189#[cfg(feature = "otlp")]
190pub fn build_otlp_layer<S>(
191    service_name: &str,
192    endpoint: &str,
193) -> Result<Box<dyn tracing_subscriber::Layer<S> + Send + Sync>, TelemetryError>
194where
195    S: tracing::Subscriber
196        + for<'span> tracing_subscriber::registry::LookupSpan<'span>
197        + Send
198        + Sync,
199{
200    use opentelemetry::trace::TracerProvider;
201    use opentelemetry_otlp::WithExportConfig;
202    use tracing_opentelemetry::OpenTelemetryLayer;
203
204    let resource = opentelemetry_sdk::Resource::builder_empty()
205        .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.to_string())])
206        .build();
207
208    // Build OTLP span exporter
209    let span_exporter = opentelemetry_otlp::SpanExporter::builder()
210        .with_tonic()
211        .with_endpoint(endpoint)
212        .build()
213        .map_err(|e| TelemetryError::Init(format!("failed to build OTLP span exporter: {e}")))?;
214
215    // Build tracer provider with batch exporter
216    let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
217        .with_batch_exporter(span_exporter)
218        .with_resource(resource.clone())
219        .build();
220
221    let tracer = tracer_provider.tracer("adk-telemetry");
222    opentelemetry::global::set_tracer_provider(tracer_provider);
223
224    // Build OTLP metric exporter
225    let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
226        .with_tonic()
227        .with_endpoint(endpoint)
228        .build()
229        .map_err(|e| TelemetryError::Init(format!("failed to build OTLP metric exporter: {e}")))?;
230
231    let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
232        .with_periodic_exporter(metric_exporter)
233        .with_resource(resource)
234        .build();
235
236    opentelemetry::global::set_meter_provider(meter_provider);
237
238    Ok(Box::new(OpenTelemetryLayer::new(tracer)))
239}
240
241/// Shutdown telemetry and flush any pending spans.
242///
243/// Should be called before application exit to ensure all telemetry data is sent.
244/// In OTel 0.28+, the tracer provider is shut down when the last reference is dropped.
245/// This function is kept for backward compatibility and explicitly drops the global provider.
246pub fn shutdown_telemetry() {
247    #[cfg(feature = "otlp")]
248    {
249        // In OTel 0.28, shutdown_tracer_provider() was removed.
250        // The SdkTracerProvider shuts down automatically when the last reference is dropped.
251        // We trigger this by replacing the global provider with a no-op, which drops the old one.
252        opentelemetry::global::set_tracer_provider(
253            opentelemetry::trace::noop::NoopTracerProvider::new(),
254        );
255    }
256}
257
258/// Initialize telemetry with ADK span exporter.
259///
260/// Creates a shared span exporter that can be used by both telemetry and the debug API.
261/// Returns the exporter so it can be passed to the debug controller.
262pub fn init_with_adk_exporter(service_name: &str) -> Result<Arc<AdkSpanExporter>, TelemetryError> {
263    let exporter = Arc::new(AdkSpanExporter::new());
264    let exporter_clone = exporter.clone();
265
266    INIT.call_once(|| {
267        let filter = EnvFilter::try_from_default_env()
268            .or_else(|_| EnvFilter::try_new("info"))
269            .unwrap_or_else(|_| EnvFilter::new("info"));
270
271        let adk_layer = AdkSpanLayer::new(exporter_clone);
272
273        tracing_subscriber::registry()
274            .with(filter)
275            .with(
276                tracing_subscriber::fmt::layer()
277                    .with_target(true)
278                    .with_thread_ids(true)
279                    .with_line_number(true),
280            )
281            .with(adk_layer)
282            .init();
283
284        tracing::info!(service.name = service_name, "telemetry initialized with ADK span exporter");
285    });
286
287    Ok(exporter)
288}