harn_vm/stdlib/
tracing.rs

1use std::collections::BTreeMap;
2use std::rc::Rc;
3use std::sync::atomic::Ordering;
4
5use crate::value::{VmError, VmValue};
6use crate::vm::Vm;
7
8use super::logging::{vm_build_log_line, VmTraceContext, VM_MIN_LOG_LEVEL, VM_TRACE_STACK};
9
10/// Finish a span started by `trace_start`: computes the elapsed duration,
11/// pops the span from the thread-local trace stack if it is on top, and
12/// returns `(name, trace_id, span_id, duration_ms)` suitable for both the
13/// default `out`-buffered `trace_end` and the bridge-streaming override
14/// registered by the ACP runner.
15pub fn finish_span_from_args(args: &[VmValue]) -> Result<(String, String, String, i64), VmError> {
16    let span = match args.first() {
17        Some(VmValue::Dict(d)) => d,
18        _ => {
19            return Err(VmError::Thrown(VmValue::String(Rc::from(
20                "trace_end: argument must be a span dict from trace_start",
21            ))));
22        }
23    };
24    let end_ms = std::time::SystemTime::now()
25        .duration_since(std::time::UNIX_EPOCH)
26        .unwrap_or_default()
27        .as_millis() as i64;
28    let start_ms = span
29        .get("start_ms")
30        .and_then(|v| v.as_int())
31        .unwrap_or(end_ms);
32    let duration_ms = end_ms - start_ms;
33    let name = span.get("name").map(|v| v.display()).unwrap_or_default();
34    let trace_id = span
35        .get("trace_id")
36        .map(|v| v.display())
37        .unwrap_or_default();
38    let span_id = span.get("span_id").map(|v| v.display()).unwrap_or_default();
39
40    VM_TRACE_STACK.with(|stack| {
41        let mut s = stack.borrow_mut();
42        if let Some(top) = s.last() {
43            if top.span_id == span_id {
44                s.pop();
45            }
46        }
47    });
48
49    Ok((name, trace_id, span_id, duration_ms))
50}
51
52pub(crate) fn current_trace_context() -> Option<(String, String)> {
53    VM_TRACE_STACK.with(|stack| {
54        stack
55            .borrow()
56            .last()
57            .map(|trace| (trace.trace_id.clone(), trace.span_id.clone()))
58    })
59}
60
61pub(crate) fn register_tracing_builtins(vm: &mut Vm) {
62    vm.register_builtin("trace_start", |args, _out| {
63        use rand::RngExt;
64        let name = args.first().map(|a| a.display()).unwrap_or_default();
65        let trace_id = VM_TRACE_STACK.with(|stack| {
66            stack
67                .borrow()
68                .last()
69                .map(|t| t.trace_id.clone())
70                .unwrap_or_else(|| {
71                    let val: u32 = rand::rng().random();
72                    format!("{val:08x}")
73                })
74        });
75        let span_id = {
76            let val: u32 = rand::rng().random();
77            format!("{val:08x}")
78        };
79        let start_ms = std::time::SystemTime::now()
80            .duration_since(std::time::UNIX_EPOCH)
81            .unwrap_or_default()
82            .as_millis() as i64;
83
84        VM_TRACE_STACK.with(|stack| {
85            stack.borrow_mut().push(VmTraceContext {
86                trace_id: trace_id.clone(),
87                span_id: span_id.clone(),
88            });
89        });
90
91        let mut span = BTreeMap::new();
92        span.insert("trace_id".to_string(), VmValue::String(Rc::from(trace_id)));
93        span.insert("span_id".to_string(), VmValue::String(Rc::from(span_id)));
94        span.insert("name".to_string(), VmValue::String(Rc::from(name)));
95        span.insert("start_ms".to_string(), VmValue::Int(start_ms));
96        Ok(VmValue::Dict(Rc::new(span)))
97    });
98
99    vm.register_builtin("trace_end", |args, out| {
100        let (name, trace_id, span_id, duration_ms) = finish_span_from_args(args)?;
101        let level_num = 1_u8;
102        if level_num >= VM_MIN_LOG_LEVEL.load(Ordering::Relaxed) {
103            let mut fields = BTreeMap::new();
104            fields.insert("trace_id".to_string(), VmValue::String(Rc::from(trace_id)));
105            fields.insert("span_id".to_string(), VmValue::String(Rc::from(span_id)));
106            fields.insert("name".to_string(), VmValue::String(Rc::from(name)));
107            fields.insert("duration_ms".to_string(), VmValue::Int(duration_ms));
108            let line = vm_build_log_line("info", "span_end", Some(&fields));
109            out.push_str(&line);
110        }
111        Ok(VmValue::Nil)
112    });
113
114    vm.register_builtin("trace_id", |_args, _out| {
115        match current_trace_context().map(|context| context.0) {
116            Some(trace_id) => Ok(VmValue::String(Rc::from(trace_id))),
117            None => Ok(VmValue::Nil),
118        }
119    });
120
121    vm.register_builtin("llm_info", |_args, _out| {
122        let raw_model = std::env::var("HARN_LLM_MODEL").unwrap_or_default();
123        let resolved = crate::llm_config::resolve_model_info(&raw_model);
124        let provider = std::env::var("HARN_LLM_PROVIDER").unwrap_or(resolved.provider);
125        let model = if raw_model.is_empty() {
126            String::new()
127        } else {
128            resolved.id
129        };
130        let api_key_set = crate::llm_config::provider_key_available(&provider);
131        let mut info = BTreeMap::new();
132        info.insert("provider".to_string(), VmValue::String(Rc::from(provider)));
133        info.insert("model".to_string(), VmValue::String(Rc::from(model)));
134        info.insert("api_key_set".to_string(), VmValue::Bool(api_key_set));
135        Ok(VmValue::Dict(Rc::new(info)))
136    });
137
138    vm.register_builtin("enable_tracing", |args, _out| {
139        let enabled = match args.first() {
140            Some(VmValue::Bool(b)) => *b,
141            _ => true,
142        };
143        crate::tracing::set_tracing_enabled(enabled);
144        Ok(VmValue::Nil)
145    });
146
147    vm.register_builtin("trace_spans", |_args, _out| {
148        let spans = crate::tracing::peek_spans();
149        let vm_spans: Vec<VmValue> = spans.iter().map(crate::tracing::span_to_vm_value).collect();
150        Ok(VmValue::List(Rc::new(vm_spans)))
151    });
152
153    vm.register_builtin("trace_summary", |_args, _out| {
154        Ok(VmValue::String(Rc::from(crate::tracing::format_summary())))
155    });
156
157    vm.register_builtin("llm_usage", |_args, _out| {
158        let (total_input, total_output, total_duration, call_count) =
159            crate::llm::peek_trace_summary();
160        let mut usage = BTreeMap::new();
161        usage.insert("input_tokens".to_string(), VmValue::Int(total_input));
162        usage.insert("output_tokens".to_string(), VmValue::Int(total_output));
163        usage.insert(
164            "total_duration_ms".to_string(),
165            VmValue::Int(total_duration),
166        );
167        usage.insert("call_count".to_string(), VmValue::Int(call_count));
168        usage.insert("total_calls".to_string(), VmValue::Int(call_count));
169        Ok(VmValue::Dict(Rc::new(usage)))
170    });
171}
harn_vm/stdlib/tracing.rs

harn_vm/stdlib/
tracing.rs
