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