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            .expect("Failed to create env filter");
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_otlp::WithExportConfig;
66    use tracing_opentelemetry::OpenTelemetryLayer;
67
68    INIT.call_once(|| {
69        // install_batch returns a Tracer directly
70        let tracer = opentelemetry_otlp::new_pipeline()
71            .tracing()
72            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
73            .with_trace_config(opentelemetry_sdk::trace::config().with_resource(
74                opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
75                    "service.name",
76                    service_name.to_string(),
77                )]),
78            ))
79            .install_batch(opentelemetry_sdk::runtime::Tokio)
80            .expect("Failed to install OTLP pipeline");
81
82        // Initialize metrics
83        let meter_provider = opentelemetry_otlp::new_pipeline()
84            .metrics(opentelemetry_sdk::runtime::Tokio)
85            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
86            .with_resource(opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
87                "service.name",
88                service_name.to_string(),
89            )]))
90            .build()
91            .expect("Failed to build meter provider");
92
93        opentelemetry::global::set_meter_provider(meter_provider);
94
95        let telemetry_layer = OpenTelemetryLayer::new(tracer);
96
97        let filter = EnvFilter::try_from_default_env()
98            .or_else(|_| EnvFilter::try_new("info"))
99            .expect("Failed to create env filter");
100
101        tracing_subscriber::registry()
102            .with(filter)
103            .with(
104                tracing_subscriber::fmt::layer()
105                    .with_target(true)
106                    .with_thread_ids(true)
107                    .with_line_number(true),
108            )
109            .with(telemetry_layer)
110            .init();
111
112        tracing::info!(
113            service.name = service_name,
114            otlp.endpoint = endpoint,
115            "telemetry initialized with OpenTelemetry"
116        );
117    });
118
119    Ok(())
120}
121
122/// Shutdown telemetry and flush any pending spans.
123///
124/// Should be called before application exit to ensure all telemetry data is sent.
125pub fn shutdown_telemetry() {
126    opentelemetry::global::shutdown_tracer_provider();
127}
128
129/// Initialize telemetry with ADK span exporter.
130///
131/// Creates a shared span exporter that can be used by both telemetry and the debug API.
132/// Returns the exporter so it can be passed to the debug controller.
133pub fn init_with_adk_exporter(service_name: &str) -> Result<Arc<AdkSpanExporter>, TelemetryError> {
134    let exporter = Arc::new(AdkSpanExporter::new());
135    let exporter_clone = exporter.clone();
136
137    INIT.call_once(|| {
138        let filter = EnvFilter::try_from_default_env()
139            .or_else(|_| EnvFilter::try_new("info"))
140            .expect("Failed to create env filter");
141
142        let adk_layer = AdkSpanLayer::new(exporter_clone);
143
144        tracing_subscriber::registry()
145            .with(filter)
146            .with(
147                tracing_subscriber::fmt::layer()
148                    .with_target(true)
149                    .with_thread_ids(true)
150                    .with_line_number(true),
151            )
152            .with(adk_layer)
153            .init();
154
155        tracing::info!(service.name = service_name, "telemetry initialized with ADK span exporter");
156    });
157
158    Ok(exporter)
159}
adk_telemetry/init.rs

adk_telemetry/
init.rs
