adk_telemetry/
span_exporter.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock};
3use tracing::{Id, Subscriber, debug};
4use tracing_subscriber::{Layer, layer::Context, registry::LookupSpan};
5
6/// ADK-Go style span exporter that stores spans by event_id
7/// Follows the pattern from APIServerSpanExporter in ADK-Go
8#[derive(Debug, Clone, Default)]
9pub struct AdkSpanExporter {
10    /// Map of event_id -> span attributes (following ADK-Go pattern)
11    trace_dict: Arc<RwLock<HashMap<String, HashMap<String, String>>>>,
12}
13
14impl AdkSpanExporter {
15    pub fn new() -> Self {
16        Self { trace_dict: Arc::new(RwLock::new(HashMap::new())) }
17    }
18
19    /// Get trace dict (following ADK-Go GetTraceDict method)
20    pub fn get_trace_dict(&self) -> HashMap<String, HashMap<String, String>> {
21        self.trace_dict.read().unwrap_or_else(|e| e.into_inner()).clone()
22    }
23
24    /// Get trace by event_id (following ADK-Go pattern)
25    pub fn get_trace_by_event_id(&self, event_id: &str) -> Option<HashMap<String, String>> {
26        debug!("AdkSpanExporter::get_trace_by_event_id called with event_id: {}", event_id);
27        let trace_dict = self.trace_dict.read().unwrap_or_else(|e| e.into_inner());
28        let result = trace_dict.get(event_id).cloned();
29        debug!("get_trace_by_event_id result for event_id '{}': {:?}", event_id, result.is_some());
30        result
31    }
32
33    /// Get all spans for a session (by filtering spans that have matching session_id)
34    pub fn get_session_trace(&self, session_id: &str) -> Vec<HashMap<String, String>> {
35        debug!("AdkSpanExporter::get_session_trace called with session_id: {}", session_id);
36        let trace_dict = self.trace_dict.read().unwrap_or_else(|e| e.into_inner());
37
38        let mut spans = Vec::new();
39        for (_event_id, attributes) in trace_dict.iter() {
40            // Check if this span belongs to the session
41            if let Some(span_session_id) = attributes.get("gcp.vertex.agent.session_id") {
42                if span_session_id == session_id {
43                    spans.push(attributes.clone());
44                }
45            }
46        }
47
48        debug!("get_session_trace result for session_id '{}': {} spans", session_id, spans.len());
49        spans
50    }
51
52    /// Internal method to store span (following ADK-Go ExportSpans pattern)
53    fn export_span(&self, span_name: &str, attributes: HashMap<String, String>) {
54        // Only capture specific span names (following ADK-Go pattern)
55        if span_name == "agent.execute"
56            || span_name == "call_llm"
57            || span_name == "send_data"
58            || span_name.starts_with("execute_tool")
59        {
60            if let Some(event_id) = attributes.get("gcp.vertex.agent.event_id") {
61                debug!(
62                    "AdkSpanExporter: Storing span '{}' with event_id '{}'",
63                    span_name, event_id
64                );
65                let mut trace_dict = self.trace_dict.write().unwrap_or_else(|e| e.into_inner());
66                trace_dict.insert(event_id.clone(), attributes);
67                debug!("AdkSpanExporter: Span stored, total event_ids: {}", trace_dict.len());
68            } else {
69                debug!("AdkSpanExporter: Skipping span '{}' - no event_id found", span_name);
70            }
71        } else {
72            debug!("AdkSpanExporter: Skipping span '{}' - not in allowed list", span_name);
73        }
74    }
75}
76
77/// Tracing layer that captures spans and exports them via AdkSpanExporter
78pub struct AdkSpanLayer {
79    exporter: Arc<AdkSpanExporter>,
80}
81
82impl AdkSpanLayer {
83    pub fn new(exporter: Arc<AdkSpanExporter>) -> Self {
84        Self { exporter }
85    }
86}
87
88#[derive(Clone)]
89struct SpanFields(HashMap<String, String>);
90
91#[derive(Clone)]
92struct SpanTiming {
93    start_time: std::time::Instant,
94}
95
96impl<S> Layer<S> for AdkSpanLayer
97where
98    S: Subscriber + for<'a> LookupSpan<'a>,
99{
100    fn on_new_span(&self, attrs: &tracing::span::Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
101        let Some(span) = ctx.span(id) else { return };
102        let mut extensions = span.extensions_mut();
103
104        // Record start time
105        extensions.insert(SpanTiming { start_time: std::time::Instant::now() });
106
107        // Capture fields
108        let mut visitor = StringVisitor::default();
109        attrs.record(&mut visitor);
110        let mut fields_map = visitor.0;
111
112        // Propagate fields from parent span (for context inheritance)
113        if let Some(parent) = span.parent() {
114            if let Some(parent_fields) = parent.extensions().get::<SpanFields>() {
115                let context_keys = [
116                    "gcp.vertex.agent.session_id",
117                    "gcp.vertex.agent.invocation_id",
118                    "gcp.vertex.agent.event_id",
119                    "gen_ai.conversation.id",
120                ];
121
122                for key in context_keys {
123                    if !fields_map.contains_key(key) {
124                        if let Some(val) = parent_fields.0.get(key) {
125                            fields_map.insert(key.to_string(), val.clone());
126                        }
127                    }
128                }
129            }
130        }
131
132        extensions.insert(SpanFields(fields_map));
133    }
134
135    fn on_record(&self, id: &Id, values: &tracing::span::Record<'_>, ctx: Context<'_, S>) {
136        let Some(span) = ctx.span(id) else { return };
137        let mut extensions = span.extensions_mut();
138        if let Some(fields) = extensions.get_mut::<SpanFields>() {
139            let mut visitor = StringVisitor::default();
140            values.record(&mut visitor);
141            for (k, v) in visitor.0 {
142                fields.0.insert(k, v);
143            }
144        }
145    }
146
147    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
148        let Some(span) = ctx.span(&id) else { return };
149        let extensions = span.extensions();
150
151        // Calculate actual duration
152        let timing = extensions.get::<SpanTiming>();
153        let end_time = std::time::Instant::now();
154        let duration_nanos =
155            timing.map(|t| end_time.duration_since(t.start_time).as_nanos() as u64).unwrap_or(0);
156
157        // Get captured fields
158        let mut attributes =
159            extensions.get::<SpanFields>().map(|f| f.0.clone()).unwrap_or_default();
160
161        // Get span name - prefer otel.name attribute (for dynamic names), fallback to metadata
162        let metadata = span.metadata();
163        let span_name =
164            attributes.get("otel.name").cloned().unwrap_or_else(|| metadata.name().to_string());
165
166        // Add span metadata and actual timing with unique IDs
167        let now_nanos = std::time::SystemTime::now()
168            .duration_since(std::time::UNIX_EPOCH)
169            .unwrap_or_default()
170            .as_nanos() as u64;
171
172        // Use invocation_id as trace_id (for grouping in UI)
173        // Use event_id as span_id (for uniqueness)
174        let invocation_id = attributes
175            .get("gcp.vertex.agent.invocation_id")
176            .cloned()
177            .unwrap_or_else(|| format!("{:016x}", id.into_u64()));
178        let event_id = attributes
179            .get("gcp.vertex.agent.event_id")
180            .cloned()
181            .unwrap_or_else(|| format!("{:016x}", id.into_u64()));
182
183        attributes.insert("span_name".to_string(), span_name.clone());
184        attributes.insert("trace_id".to_string(), invocation_id); // Group by invocation
185        attributes.insert("span_id".to_string(), event_id); // Unique per span
186        attributes.insert("start_time".to_string(), (now_nanos - duration_nanos).to_string());
187        attributes.insert("end_time".to_string(), now_nanos.to_string());
188
189        // Don't set parent_span_id to keep all spans at same level like ADK-Go
190
191        // Export the span
192        self.exporter.export_span(&span_name, attributes);
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use std::sync::Arc;
200    use tracing_subscriber::layer::SubscriberExt;
201
202    #[test]
203    fn test_conversation_id_propagates_to_child_spans() {
204        let exporter = Arc::new(AdkSpanExporter::new());
205        let layer = AdkSpanLayer::new(exporter.clone());
206        let subscriber = tracing_subscriber::registry().with(layer);
207
208        tracing::subscriber::with_default(subscriber, || {
209            let parent = tracing::info_span!(
210                "agent.execute",
211                "gcp.vertex.agent.event_id" = "evt-parent",
212                "gcp.vertex.agent.invocation_id" = "inv-1",
213                "gcp.vertex.agent.session_id" = "session-1",
214                "gen_ai.conversation.id" = "session-1",
215                "agent.name" = "test-agent"
216            );
217
218            let _parent_guard = parent.enter();
219
220            let child = tracing::info_span!(
221                "call_llm",
222                "gcp.vertex.agent.event_id" = "evt-child",
223                "gcp.vertex.agent.llm_request" = "{}"
224            );
225            let _child_guard = child.enter();
226            tracing::info!("child span body");
227        });
228
229        let child_trace =
230            exporter.get_trace_by_event_id("evt-child").expect("child span should be exported");
231        assert_eq!(
232            child_trace.get("gen_ai.conversation.id").map(String::as_str),
233            Some("session-1")
234        );
235    }
236}
237
238#[derive(Default)]
239struct StringVisitor(HashMap<String, String>);
240
241impl tracing::field::Visit for StringVisitor {
242    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
243        self.0.insert(field.name().to_string(), format!("{:?}", value));
244    }
245
246    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
247        self.0.insert(field.name().to_string(), value.to_string());
248    }
249
250    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
251        self.0.insert(field.name().to_string(), value.to_string());
252    }
253
254    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
255        self.0.insert(field.name().to_string(), value.to_string());
256    }
257
258    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
259        self.0.insert(field.name().to_string(), value.to_string());
260    }
261
262    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
263        self.0.insert(field.name().to_string(), value.to_string());
264    }
265}
adk_telemetry/span_exporter.rs

adk_telemetry/
span_exporter.rs
