harn_vm/
tracing.rs

1//! Pipeline Observability: structured tracing spans with parent/child relationships.
2//!
3//! When tracing is enabled (`vm.enable_tracing()`), the VM automatically emits
4//! spans for pipeline execution, function calls, LLM calls, tool invocations,
5//! imports, and async operations. Spans form a tree via parent_span_id.
6//!
7//! Access via builtins: `trace_spans()` returns all completed spans,
8//! `trace_summary()` returns a formatted summary.
9
10use std::cell::RefCell;
11use std::collections::BTreeMap;
12use std::rc::Rc;
13use std::time::Instant;
14
15use crate::value::VmValue;
16
17// =============================================================================
18// Span types
19// =============================================================================
20
21/// The kind of operation a span represents.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum SpanKind {
24    Pipeline,
25    FnCall,
26    LlmCall,
27    ToolCall,
28    Import,
29    Parallel,
30    Spawn,
31}
32
33impl SpanKind {
34    pub fn as_str(self) -> &'static str {
35        match self {
36            Self::Pipeline => "pipeline",
37            Self::FnCall => "fn_call",
38            Self::LlmCall => "llm_call",
39            Self::ToolCall => "tool_call",
40            Self::Import => "import",
41            Self::Parallel => "parallel",
42            Self::Spawn => "spawn",
43        }
44    }
45}
46
47/// A completed tracing span.
48#[derive(Debug, Clone)]
49pub struct Span {
50    pub span_id: u64,
51    pub parent_id: Option<u64>,
52    pub kind: SpanKind,
53    pub name: String,
54    pub start_ms: u64,
55    pub duration_ms: u64,
56    pub metadata: BTreeMap<String, serde_json::Value>,
57}
58
59/// An in-flight span (not yet completed).
60struct OpenSpan {
61    span_id: u64,
62    parent_id: Option<u64>,
63    kind: SpanKind,
64    name: String,
65    started_at: Instant,
66    metadata: BTreeMap<String, serde_json::Value>,
67}
68
69// =============================================================================
70// Collector
71// =============================================================================
72
73/// Thread-local span collector. Accumulates completed spans and tracks the
74/// active span stack for automatic parent assignment.
75pub struct SpanCollector {
76    next_id: u64,
77    /// Stack of open span IDs — the top is the current active span.
78    active_stack: Vec<u64>,
79    /// Open (in-flight) spans keyed by ID.
80    open: BTreeMap<u64, OpenSpan>,
81    /// Completed spans in chronological order.
82    completed: Vec<Span>,
83    /// Epoch for relative timing.
84    epoch: Instant,
85}
86
87impl Default for SpanCollector {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93impl SpanCollector {
94    pub fn new() -> Self {
95        Self {
96            next_id: 1,
97            active_stack: Vec::new(),
98            open: BTreeMap::new(),
99            completed: Vec::new(),
100            epoch: Instant::now(),
101        }
102    }
103
104    /// Start a new span. Returns the span ID.
105    pub fn start(&mut self, kind: SpanKind, name: String) -> u64 {
106        let id = self.next_id;
107        self.next_id += 1;
108        let parent_id = self.active_stack.last().copied();
109        let now = Instant::now();
110
111        // Emit structured span-start event to registered sinks.
112        crate::events::emit_span_start(id, parent_id, &name, kind.as_str(), BTreeMap::new());
113
114        self.open.insert(
115            id,
116            OpenSpan {
117                span_id: id,
118                parent_id,
119                kind,
120                name,
121                started_at: now,
122                metadata: BTreeMap::new(),
123            },
124        );
125        self.active_stack.push(id);
126        id
127    }
128
129    /// Attach metadata to an open span.
130    pub fn set_metadata(&mut self, span_id: u64, key: &str, value: serde_json::Value) {
131        if let Some(span) = self.open.get_mut(&span_id) {
132            span.metadata.insert(key.to_string(), value);
133        }
134    }
135
136    /// End a span. Moves it from open to completed.
137    pub fn end(&mut self, span_id: u64) {
138        if let Some(span) = self.open.remove(&span_id) {
139            let duration = span.started_at.elapsed();
140            let start_ms = span.started_at.duration_since(self.epoch).as_millis() as u64;
141            let duration_ms = duration.as_millis() as u64;
142
143            // Emit structured span-end event to registered sinks.
144            let mut end_meta = span.metadata.clone();
145            end_meta.insert(
146                "duration_ms".to_string(),
147                serde_json::Value::Number(serde_json::Number::from(duration_ms)),
148            );
149            crate::events::emit_span_end(span_id, end_meta);
150
151            self.completed.push(Span {
152                span_id: span.span_id,
153                parent_id: span.parent_id,
154                kind: span.kind,
155                name: span.name,
156                start_ms,
157                duration_ms,
158                metadata: span.metadata,
159            });
160
161            // Remove from active stack
162            if let Some(pos) = self.active_stack.iter().rposition(|&id| id == span_id) {
163                self.active_stack.remove(pos);
164            }
165        }
166    }
167
168    /// Get the current active span ID (if any).
169    pub fn current_span_id(&self) -> Option<u64> {
170        self.active_stack.last().copied()
171    }
172
173    /// Take all completed spans (drains the collector).
174    pub fn take_spans(&mut self) -> Vec<Span> {
175        std::mem::take(&mut self.completed)
176    }
177
178    /// Peek at all completed spans (non-destructive).
179    pub fn spans(&self) -> &[Span] {
180        &self.completed
181    }
182
183    /// Reset the collector.
184    pub fn reset(&mut self) {
185        self.active_stack.clear();
186        self.open.clear();
187        self.completed.clear();
188        self.next_id = 1;
189        self.epoch = Instant::now();
190    }
191}
192
193// =============================================================================
194// Thread-local collector
195// =============================================================================
196
197thread_local! {
198    static COLLECTOR: RefCell<SpanCollector> = RefCell::new(SpanCollector::new());
199    static TRACING_ENABLED: RefCell<bool> = const { RefCell::new(false) };
200}
201
202/// Enable or disable VM tracing for the current thread.
203pub fn set_tracing_enabled(enabled: bool) {
204    TRACING_ENABLED.with(|e| *e.borrow_mut() = enabled);
205    if enabled {
206        COLLECTOR.with(|c| c.borrow_mut().reset());
207    }
208}
209
210/// Check if tracing is enabled.
211pub fn is_tracing_enabled() -> bool {
212    TRACING_ENABLED.with(|e| *e.borrow())
213}
214
215/// Start a span (no-op if tracing disabled). Returns span ID or 0.
216pub fn span_start(kind: SpanKind, name: String) -> u64 {
217    if !is_tracing_enabled() {
218        return 0;
219    }
220    COLLECTOR.with(|c| c.borrow_mut().start(kind, name))
221}
222
223/// Attach metadata to an open span (no-op if span_id is 0).
224pub fn span_set_metadata(span_id: u64, key: &str, value: serde_json::Value) {
225    if span_id == 0 {
226        return;
227    }
228    COLLECTOR.with(|c| c.borrow_mut().set_metadata(span_id, key, value));
229}
230
231/// End a span (no-op if span_id is 0).
232pub fn span_end(span_id: u64) {
233    if span_id == 0 {
234        return;
235    }
236    COLLECTOR.with(|c| c.borrow_mut().end(span_id));
237}
238
239/// Get the currently active span id, if tracing is enabled and a span is open.
240pub fn current_span_id() -> Option<u64> {
241    if !is_tracing_enabled() {
242        return None;
243    }
244    COLLECTOR.with(|c| c.borrow().current_span_id())
245}
246
247/// Take all completed spans.
248pub fn take_spans() -> Vec<Span> {
249    COLLECTOR.with(|c| c.borrow_mut().take_spans())
250}
251
252/// Peek at completed spans (cloned).
253pub fn peek_spans() -> Vec<Span> {
254    COLLECTOR.with(|c| c.borrow().spans().to_vec())
255}
256
257/// Reset the tracing collector.
258pub fn reset_tracing() {
259    COLLECTOR.with(|c| c.borrow_mut().reset());
260}
261
262// =============================================================================
263// VmValue conversion
264// =============================================================================
265
266/// Convert a span to a VmValue dict for user access.
267pub fn span_to_vm_value(span: &Span) -> VmValue {
268    let mut d = BTreeMap::new();
269    d.insert("span_id".into(), VmValue::Int(span.span_id as i64));
270    d.insert(
271        "parent_id".into(),
272        span.parent_id
273            .map(|id| VmValue::Int(id as i64))
274            .unwrap_or(VmValue::Nil),
275    );
276    d.insert("kind".into(), VmValue::String(Rc::from(span.kind.as_str())));
277    d.insert("name".into(), VmValue::String(Rc::from(span.name.as_str())));
278    d.insert("start_ms".into(), VmValue::Int(span.start_ms as i64));
279    d.insert("duration_ms".into(), VmValue::Int(span.duration_ms as i64));
280
281    if !span.metadata.is_empty() {
282        let meta: BTreeMap<String, VmValue> = span
283            .metadata
284            .iter()
285            .map(|(k, v)| (k.clone(), crate::stdlib::json_to_vm_value(v)))
286            .collect();
287        d.insert("metadata".into(), VmValue::Dict(Rc::new(meta)));
288    }
289
290    VmValue::Dict(Rc::new(d))
291}
292
293/// Generate a formatted summary of all spans.
294pub fn format_summary() -> String {
295    let spans = peek_spans();
296    if spans.is_empty() {
297        return "No spans recorded.".into();
298    }
299
300    let mut lines = Vec::new();
301    let total_ms: u64 = spans
302        .iter()
303        .filter(|s| s.parent_id.is_none())
304        .map(|s| s.duration_ms)
305        .sum();
306
307    lines.push(format!("Trace: {} spans, {total_ms}ms total", spans.len()));
308    lines.push(String::new());
309
310    // Build tree structure
311    fn print_tree(spans: &[Span], parent_id: Option<u64>, depth: usize, lines: &mut Vec<String>) {
312        let children: Vec<&Span> = spans.iter().filter(|s| s.parent_id == parent_id).collect();
313        for span in children {
314            let indent = "  ".repeat(depth);
315            let meta_str = if span.metadata.is_empty() {
316                String::new()
317            } else {
318                let parts: Vec<String> = span
319                    .metadata
320                    .iter()
321                    .map(|(k, v)| format!("{k}={v}"))
322                    .collect();
323                format!(" ({})", parts.join(", "))
324            };
325            lines.push(format!(
326                "{indent}{} {} {}ms{meta_str}",
327                span.kind.as_str(),
328                span.name,
329                span.duration_ms,
330            ));
331            print_tree(spans, Some(span.span_id), depth + 1, lines);
332        }
333    }
334
335    print_tree(&spans, None, 0, &mut lines);
336    lines.join("\n")
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn test_span_collector_basic() {
345        let mut c = SpanCollector::new();
346        let id = c.start(SpanKind::Pipeline, "main".into());
347        assert_eq!(id, 1);
348        assert_eq!(c.current_span_id(), Some(1));
349        c.end(id);
350        assert_eq!(c.current_span_id(), None);
351        assert_eq!(c.spans().len(), 1);
352        assert_eq!(c.spans()[0].name, "main");
353        assert_eq!(c.spans()[0].parent_id, None);
354    }
355
356    #[test]
357    fn test_span_parent_child() {
358        let mut c = SpanCollector::new();
359        let parent = c.start(SpanKind::Pipeline, "main".into());
360        let child = c.start(SpanKind::FnCall, "helper".into());
361        c.end(child);
362        c.end(parent);
363        assert_eq!(c.spans().len(), 2);
364        assert_eq!(c.spans()[0].parent_id, Some(parent));
365        assert_eq!(c.spans()[1].parent_id, None);
366    }
367
368    #[test]
369    fn test_span_metadata() {
370        let mut c = SpanCollector::new();
371        let id = c.start(SpanKind::LlmCall, "gpt-4".into());
372        c.set_metadata(id, "tokens", serde_json::json!(100));
373        c.end(id);
374        assert_eq!(c.spans()[0].metadata["tokens"], serde_json::json!(100));
375    }
376
377    #[test]
378    fn test_noop_when_disabled() {
379        set_tracing_enabled(false);
380        let id = span_start(SpanKind::Pipeline, "test".into());
381        assert_eq!(id, 0);
382        span_end(id); // should not panic
383    }
384}
harn_vm/tracing.rs

harn_vm/
tracing.rs
