harn_vm/llm/
trace.rs

1use std::cell::RefCell;
2
3// =============================================================================
4// LLM trace log (thread-local for async-safe access)
5// =============================================================================
6
7/// A single LLM call trace entry.
8#[derive(Debug, Clone)]
9pub struct LlmTraceEntry {
10    pub model: String,
11    pub input_tokens: i64,
12    pub output_tokens: i64,
13    pub duration_ms: u64,
14}
15
16thread_local! {
17    static LLM_TRACE: RefCell<Vec<LlmTraceEntry>> = const { RefCell::new(Vec::new()) };
18    static LLM_TRACING_ENABLED: RefCell<bool> = const { RefCell::new(false) };
19}
20
21/// Enable LLM tracing for the current thread.
22pub fn enable_tracing() {
23    LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = true);
24}
25
26/// Get and clear the trace log.
27pub fn take_trace() -> Vec<LlmTraceEntry> {
28    LLM_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
29}
30
31/// Clone the current trace log without consuming it.
32pub fn peek_trace() -> Vec<LlmTraceEntry> {
33    LLM_TRACE.with(|v| v.borrow().clone())
34}
35
36/// Summarize trace usage without consuming entries.
37pub fn peek_trace_summary() -> (i64, i64, i64, i64) {
38    LLM_TRACE.with(|v| {
39        let entries = v.borrow();
40        let mut input = 0i64;
41        let mut output = 0i64;
42        let mut duration = 0i64;
43        let count = entries.len() as i64;
44        for e in entries.iter() {
45            input += e.input_tokens;
46            output += e.output_tokens;
47            duration += e.duration_ms as i64;
48        }
49        (input, output, duration, count)
50    })
51}
52
53/// Reset thread-local trace state. Call between test runs.
54pub(crate) fn reset_trace_state() {
55    LLM_TRACE.with(|v| v.borrow_mut().clear());
56    LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = false);
57}
58
59pub(crate) fn trace_llm_call(entry: LlmTraceEntry) {
60    LLM_TRACING_ENABLED.with(|enabled| {
61        if *enabled.borrow() {
62            LLM_TRACE.with(|v| v.borrow_mut().push(entry));
63        }
64    });
65}
harn_vm/llm/trace.rs

harn_vm/llm/
trace.rs
