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

harn_vm/
tracing.rs
