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

harn_vm/
tracing.rs
