harn_vm/stdlib/
tracing.rs

1use crate::value::VmDictExt;
2use std::sync::atomic::Ordering;
3
4use crate::stdlib::macros::{harn_builtin, VmBuiltinDef};
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(arcstr::ArcStr::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    for def in MODULE_BUILTINS {
63        vm.register_builtin_def(def);
64    }
65}
66
67#[harn_builtin(sig = "trace_start(...args: any) -> dict", category = "tracing")]
68fn trace_start_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
69    use rand::RngExt;
70    let name = args.first().map(|a| a.display()).unwrap_or_default();
71    let trace_id = VM_TRACE_STACK.with(|stack| {
72        stack
73            .borrow()
74            .last()
75            .map(|t| t.trace_id.clone())
76            .unwrap_or_else(|| {
77                let val: u32 = rand::rng().random();
78                format!("{val:08x}")
79            })
80    });
81    let span_id = {
82        let val: u32 = rand::rng().random();
83        format!("{val:08x}")
84    };
85    let start_ms = std::time::SystemTime::now()
86        .duration_since(std::time::UNIX_EPOCH)
87        .unwrap_or_default()
88        .as_millis() as i64;
89
90    VM_TRACE_STACK.with(|stack| {
91        stack.borrow_mut().push(VmTraceContext {
92            trace_id: trace_id.clone(),
93            span_id: span_id.clone(),
94        });
95    });
96
97    let mut span = crate::value::DictMap::new();
98    span.put_str("trace_id", trace_id);
99    span.put_str("span_id", span_id);
100    span.put_str("name", name);
101    span.insert(crate::value::intern_key("start_ms"), VmValue::Int(start_ms));
102    Ok(VmValue::dict(span))
103}
104
105#[harn_builtin(sig = "trace_end(...args: any) -> nil", category = "tracing")]
106fn trace_end_impl(args: &[VmValue], out: &mut String) -> Result<VmValue, VmError> {
107    let (name, trace_id, span_id, duration_ms) = finish_span_from_args(args)?;
108    let level_num = 1_u8;
109    if level_num >= VM_MIN_LOG_LEVEL.load(Ordering::Relaxed) {
110        let mut fields = crate::value::DictMap::new();
111        fields.put_str("trace_id", trace_id);
112        fields.put_str("span_id", span_id);
113        fields.put_str("name", name);
114        fields.insert(
115            crate::value::intern_key("duration_ms"),
116            VmValue::Int(duration_ms),
117        );
118        let line = vm_build_log_line("info", "span_end", Some(&fields));
119        out.push_str(&line);
120    }
121    Ok(VmValue::Nil)
122}
123
124#[harn_builtin(sig = "trace_id(...args: any) -> string?", category = "tracing")]
125fn trace_id_impl(_args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
126    match current_trace_context().map(|context| context.0) {
127        Some(trace_id) => Ok(VmValue::String(arcstr::ArcStr::from(trace_id))),
128        None => Ok(VmValue::Nil),
129    }
130}
131
132#[harn_builtin(sig = "llm_info() -> dict", category = "tracing")]
133fn llm_info_impl(_args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
134    let raw_model = std::env::var("HARN_LLM_MODEL").unwrap_or_default();
135    let resolved = crate::llm_config::resolve_model_info(&raw_model);
136    let provider = std::env::var("HARN_LLM_PROVIDER").unwrap_or(resolved.provider);
137    let model = if raw_model.is_empty() {
138        String::new()
139    } else {
140        resolved.id
141    };
142    let api_key_set = crate::llm_config::provider_key_available(&provider);
143    let mut info = crate::value::DictMap::new();
144    info.put_str("provider", provider);
145    info.put_str("model", model);
146    info.insert(
147        crate::value::intern_key("api_key_set"),
148        VmValue::Bool(api_key_set),
149    );
150    Ok(VmValue::dict(info))
151}
152
153#[harn_builtin(sig = "enable_tracing(enabled?: bool) -> nil", category = "tracing")]
154fn enable_tracing_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
155    let enabled = match args.first() {
156        Some(VmValue::Bool(b)) => *b,
157        _ => true,
158    };
159    crate::tracing::set_tracing_enabled(enabled);
160    Ok(VmValue::Nil)
161}
162
163#[harn_builtin(sig = "trace_spans(...args: any) -> list", category = "tracing")]
164fn trace_spans_impl(_args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
165    let spans = crate::tracing::peek_spans();
166    let vm_spans: Vec<VmValue> = spans.iter().map(crate::tracing::span_to_vm_value).collect();
167    Ok(VmValue::List(std::sync::Arc::new(vm_spans)))
168}
169
170#[harn_builtin(sig = "trace_summary(...args: any) -> string", category = "tracing")]
171fn trace_summary_impl(_args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
172    Ok(VmValue::String(arcstr::ArcStr::from(
173        crate::tracing::format_summary(),
174    )))
175}
176
177#[harn_builtin(
178    sig = "__lifecycle_span_start(name: string) -> int",
179    category = "tracing"
180)]
181fn lifecycle_span_start_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
182    let name = args.first().map(|a| a.display()).unwrap_or_default();
183    let id = crate::tracing::span_start(crate::tracing::SpanKind::FnCall, name);
184    Ok(VmValue::Int(id as i64))
185}
186
187#[harn_builtin(
188    sig = "__lifecycle_span_end(span_id: int) -> nil",
189    category = "tracing"
190)]
191fn lifecycle_span_end_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
192    let span_id = match args.first().and_then(|v| v.as_int()) {
193        Some(id) => id,
194        None => return Ok(VmValue::Nil),
195    };
196    if span_id < 0 {
197        return Ok(VmValue::Nil);
198    }
199    crate::tracing::span_end(span_id as u64);
200    Ok(VmValue::Nil)
201}
202
203#[harn_builtin(sig = "llm_usage() -> dict", category = "tracing")]
204fn llm_usage_impl(_args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
205    let (total_input, total_output, total_duration, call_count) = crate::llm::peek_trace_summary();
206    let mut usage = crate::value::DictMap::new();
207    usage.insert(
208        crate::value::intern_key("input_tokens"),
209        VmValue::Int(total_input),
210    );
211    usage.insert(
212        crate::value::intern_key("output_tokens"),
213        VmValue::Int(total_output),
214    );
215    usage.insert(
216        crate::value::intern_key("total_duration_ms"),
217        VmValue::Int(total_duration),
218    );
219    usage.insert(
220        crate::value::intern_key("call_count"),
221        VmValue::Int(call_count),
222    );
223    usage.insert(
224        crate::value::intern_key("total_calls"),
225        VmValue::Int(call_count),
226    );
227    Ok(VmValue::dict(usage))
228}
229
230pub(crate) const MODULE_BUILTINS: &[&VmBuiltinDef] = &[
231    &TRACE_START_IMPL_DEF,
232    &TRACE_END_IMPL_DEF,
233    &TRACE_ID_IMPL_DEF,
234    &LLM_INFO_IMPL_DEF,
235    &ENABLE_TRACING_IMPL_DEF,
236    &TRACE_SPANS_IMPL_DEF,
237    &TRACE_SUMMARY_IMPL_DEF,
238    &LIFECYCLE_SPAN_START_IMPL_DEF,
239    &LIFECYCLE_SPAN_END_IMPL_DEF,
240    &LLM_USAGE_IMPL_DEF,
241];
harn_vm/stdlib/tracing.rs

harn_vm/stdlib/
tracing.rs
