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

harn_vm/
tracing.rs
