Skip to main content

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}
66
67// =============================================================================
68// Structured agent trace events
69// =============================================================================
70
71/// Fine-grained event emitted during agent loop execution. Captures tool
72/// calls, LLM calls, interventions, compaction, and phase changes so
73/// downstream consumers (portal, burin-code) can display execution traces
74/// without reconstructing them from raw JSON.
75#[derive(Debug, Clone, serde::Serialize)]
76#[serde(tag = "type", rename_all = "snake_case")]
77pub enum AgentTraceEvent {
78    LlmCall {
79        call_id: String,
80        model: String,
81        input_tokens: i64,
82        output_tokens: i64,
83        cache_tokens: i64,
84        duration_ms: u64,
85        iteration: usize,
86    },
87    ToolExecution {
88        tool_name: String,
89        tool_use_id: String,
90        duration_ms: u64,
91        status: String,
92        classification: String,
93        iteration: usize,
94    },
95    ToolRejected {
96        tool_name: String,
97        reason: String,
98        iteration: usize,
99    },
100    LoopIntervention {
101        tool_name: String,
102        kind: String,
103        count: usize,
104        iteration: usize,
105    },
106    ContextCompaction {
107        archived_messages: usize,
108        new_summary_len: usize,
109        iteration: usize,
110    },
111    PhaseChange {
112        from_phase: String,
113        to_phase: String,
114        iteration: usize,
115    },
116    LoopComplete {
117        status: String,
118        iterations: usize,
119        total_duration_ms: u64,
120        tools_used: Vec<String>,
121        successful_tools: Vec<String>,
122    },
123}
124
125thread_local! {
126    static AGENT_TRACE: RefCell<Vec<AgentTraceEvent>> = const { RefCell::new(Vec::new()) };
127}
128
129/// Emit an agent trace event.
130pub(crate) fn emit_agent_event(event: AgentTraceEvent) {
131    AGENT_TRACE.with(|v| v.borrow_mut().push(event));
132}
133
134/// Get and clear the agent trace log.
135pub fn take_agent_trace() -> Vec<AgentTraceEvent> {
136    AGENT_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
137}
138
139/// Clone the current agent trace log without consuming it.
140pub fn peek_agent_trace() -> Vec<AgentTraceEvent> {
141    AGENT_TRACE.with(|v| v.borrow().clone())
142}
143
144/// Produce a rolled-up summary of agent trace events as JSON.
145pub fn agent_trace_summary() -> serde_json::Value {
146    AGENT_TRACE.with(|v| {
147        let events = v.borrow();
148        let mut llm_calls = 0usize;
149        let mut tool_executions = 0usize;
150        let mut tool_rejections = 0usize;
151        let mut interventions = 0usize;
152        let mut compactions = 0usize;
153        let mut total_input_tokens = 0i64;
154        let mut total_output_tokens = 0i64;
155        let mut total_llm_duration_ms = 0u64;
156        let mut total_tool_duration_ms = 0u64;
157        let mut tools_used: Vec<String> = Vec::new();
158        let mut status = "unknown".to_string();
159        let mut iterations = 0usize;
160        let mut total_duration_ms = 0u64;
161
162        for event in events.iter() {
163            match event {
164                AgentTraceEvent::LlmCall {
165                    input_tokens,
166                    output_tokens,
167                    duration_ms,
168                    ..
169                } => {
170                    llm_calls += 1;
171                    total_input_tokens += input_tokens;
172                    total_output_tokens += output_tokens;
173                    total_llm_duration_ms += duration_ms;
174                }
175                AgentTraceEvent::ToolExecution {
176                    tool_name,
177                    duration_ms,
178                    ..
179                } => {
180                    tool_executions += 1;
181                    total_tool_duration_ms += duration_ms;
182                    if !tools_used.contains(tool_name) {
183                        tools_used.push(tool_name.clone());
184                    }
185                }
186                AgentTraceEvent::ToolRejected { .. } => {
187                    tool_rejections += 1;
188                }
189                AgentTraceEvent::LoopIntervention { .. } => {
190                    interventions += 1;
191                }
192                AgentTraceEvent::ContextCompaction { .. } => {
193                    compactions += 1;
194                }
195                AgentTraceEvent::PhaseChange { .. } => {}
196                AgentTraceEvent::LoopComplete {
197                    status: s,
198                    iterations: i,
199                    total_duration_ms: d,
200                    ..
201                } => {
202                    status = s.clone();
203                    iterations = *i;
204                    total_duration_ms = *d;
205                }
206            }
207        }
208
209        serde_json::json!({
210            "status": status,
211            "iterations": iterations,
212            "total_duration_ms": total_duration_ms,
213            "llm_calls": llm_calls,
214            "tool_executions": tool_executions,
215            "tool_rejections": tool_rejections,
216            "interventions": interventions,
217            "compactions": compactions,
218            "total_input_tokens": total_input_tokens,
219            "total_output_tokens": total_output_tokens,
220            "total_llm_duration_ms": total_llm_duration_ms,
221            "total_tool_duration_ms": total_tool_duration_ms,
222            "tools_used": tools_used,
223        })
224    })
225}
226
227/// Reset agent trace state. Call between test runs.
228pub(crate) fn reset_agent_trace_state() {
229    AGENT_TRACE.with(|v| v.borrow_mut().clear());
230}