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/// Initialize basic telemetry with console logging
11///
12/// # Arguments
13/// * `service_name` - Name of the service for trace identification
14///
15/// # Example
16/// ```
17/// use adk_telemetry::init_telemetry;
18/// init_telemetry("my-agent-service").expect("Failed to initialize telemetry");
19/// ```
20pub fn init_telemetry(service_name: &str) -> Result<(), Box<dyn std::error::Error>> {
21    INIT.call_once(|| {
22        let filter = EnvFilter::try_from_default_env()
23            .or_else(|_| EnvFilter::try_new("info"))
24            .expect("Failed to create env filter");
25
26        tracing_subscriber::registry()
27            .with(filter)
28            .with(
29                tracing_subscriber::fmt::layer()
30                    .with_target(true)
31                    .with_thread_ids(true)
32                    .with_line_number(true),
33            )
34            .init();
35
36        tracing::info!(service.name = service_name, "Telemetry initialized");
37    });
38
39    Ok(())
40}
41
42/// Initialize telemetry with OpenTelemetry OTLP export
43///
44/// Enables distributed tracing by exporting spans to an OTLP collector.
45///
46/// # Arguments
47/// * `service_name` - Name of the service for trace identification
48/// * `endpoint` - OTLP collector endpoint (e.g., "http://localhost:4317")
49///
50/// # Example
51/// ```no_run
52/// use adk_telemetry::init_with_otlp;
53/// init_with_otlp("my-agent", "http://localhost:4317")
54///     .expect("Failed to initialize telemetry");
55/// ```
56pub fn init_with_otlp(
57    service_name: &str,
58    endpoint: &str,
59) -> Result<(), Box<dyn std::error::Error>> {
60    use opentelemetry_otlp::WithExportConfig;
61    use tracing_opentelemetry::OpenTelemetryLayer;
62
63    INIT.call_once(|| {
64        // install_batch returns a Tracer directly
65        let tracer = opentelemetry_otlp::new_pipeline()
66            .tracing()
67            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
68            .with_trace_config(opentelemetry_sdk::trace::config().with_resource(
69                opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
70                    "service.name",
71                    service_name.to_string(),
72                )]),
73            ))
74            .install_batch(opentelemetry_sdk::runtime::Tokio)
75            .expect("Failed to install OTLP pipeline");
76
77        // Initialize metrics
78        let meter_provider = opentelemetry_otlp::new_pipeline()
79            .metrics(opentelemetry_sdk::runtime::Tokio)
80            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
81            .with_resource(opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
82                "service.name",
83                service_name.to_string(),
84            )]))
85            .build()
86            .expect("Failed to build meter provider");
87
88        opentelemetry::global::set_meter_provider(meter_provider);
89
90        let telemetry_layer = OpenTelemetryLayer::new(tracer);
91
92        let filter = EnvFilter::try_from_default_env()
93            .or_else(|_| EnvFilter::try_new("info"))
94            .expect("Failed to create env filter");
95
96        tracing_subscriber::registry()
97            .with(filter)
98            .with(
99                tracing_subscriber::fmt::layer()
100                    .with_target(true)
101                    .with_thread_ids(true)
102                    .with_line_number(true),
103            )
104            .with(telemetry_layer)
105            .init();
106
107        tracing::info!(
108            service.name = service_name,
109            otlp.endpoint = endpoint,
110            "Telemetry initialized with OpenTelemetry"
111        );
112    });
113
114    Ok(())
115}
116
117/// Shutdown telemetry and flush any pending spans
118///
119/// Should be called before application exit to ensure all telemetry data is sent.
120pub fn shutdown_telemetry() {
121    opentelemetry::global::shutdown_tracer_provider();
122}
123
124/// Initialize telemetry with ADK-Go style span exporter
125/// This creates a shared span exporter that can be used by both telemetry and debug API
126/// Returns the exporter so it can be passed to the debug controller
127pub fn init_with_adk_exporter(
128    service_name: &str,
129) -> Result<Arc<AdkSpanExporter>, Box<dyn std::error::Error>> {
130    let exporter = Arc::new(AdkSpanExporter::new());
131    let exporter_clone = exporter.clone();
132
133    INIT.call_once(|| {
134        let filter = EnvFilter::try_from_default_env()
135            .or_else(|_| EnvFilter::try_new("info"))
136            .expect("Failed to create env filter");
137
138        let adk_layer = AdkSpanLayer::new(exporter_clone);
139
140        tracing_subscriber::registry()
141            .with(filter)
142            .with(
143                tracing_subscriber::fmt::layer()
144                    .with_target(true)
145                    .with_thread_ids(true)
146                    .with_line_number(true),
147            )
148            .with(adk_layer)
149            .init();
150
151        tracing::info!(service.name = service_name, "Telemetry initialized with ADK span exporter");
152    });
153
154    Ok(exporter)
155}