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    /// Emitted when `llm_call` re-prompts the model after the previous
124    /// response failed `output_schema` validation. One event per retry;
125    /// `attempt` counts retries (the initial call is attempt 0 and
126    /// produces no event; the first retry emits `attempt: 1`).
127    SchemaRetry {
128        attempt: usize,
129        errors: Vec<String>,
130        nudge_used: bool,
131    },
132}
133
134thread_local! {
135    static AGENT_TRACE: RefCell<Vec<AgentTraceEvent>> = const { RefCell::new(Vec::new()) };
136}
137
138/// Emit an agent trace event.
139pub(crate) fn emit_agent_event(event: AgentTraceEvent) {
140    AGENT_TRACE.with(|v| v.borrow_mut().push(event));
141}
142
143/// Get and clear the agent trace log.
144pub fn take_agent_trace() -> Vec<AgentTraceEvent> {
145    AGENT_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
146}
147
148/// Clone the current agent trace log without consuming it.
149pub fn peek_agent_trace() -> Vec<AgentTraceEvent> {
150    AGENT_TRACE.with(|v| v.borrow().clone())
151}
152
153/// Produce a rolled-up summary of agent trace events as JSON.
154pub fn agent_trace_summary() -> serde_json::Value {
155    AGENT_TRACE.with(|v| {
156        let events = v.borrow();
157        let mut llm_calls = 0usize;
158        let mut tool_executions = 0usize;
159        let mut tool_rejections = 0usize;
160        let mut interventions = 0usize;
161        let mut compactions = 0usize;
162        let mut total_input_tokens = 0i64;
163        let mut total_output_tokens = 0i64;
164        let mut total_llm_duration_ms = 0u64;
165        let mut total_tool_duration_ms = 0u64;
166        let mut tools_used: Vec<String> = Vec::new();
167        let mut status = "unknown".to_string();
168        let mut iterations = 0usize;
169        let mut total_duration_ms = 0u64;
170
171        for event in events.iter() {
172            match event {
173                AgentTraceEvent::LlmCall {
174                    input_tokens,
175                    output_tokens,
176                    duration_ms,
177                    ..
178                } => {
179                    llm_calls += 1;
180                    total_input_tokens += input_tokens;
181                    total_output_tokens += output_tokens;
182                    total_llm_duration_ms += duration_ms;
183                }
184                AgentTraceEvent::ToolExecution {
185                    tool_name,
186                    duration_ms,
187                    ..
188                } => {
189                    tool_executions += 1;
190                    total_tool_duration_ms += duration_ms;
191                    if !tools_used.contains(tool_name) {
192                        tools_used.push(tool_name.clone());
193                    }
194                }
195                AgentTraceEvent::ToolRejected { .. } => {
196                    tool_rejections += 1;
197                }
198                AgentTraceEvent::LoopIntervention { .. } => {
199                    interventions += 1;
200                }
201                AgentTraceEvent::ContextCompaction { .. } => {
202                    compactions += 1;
203                }
204                AgentTraceEvent::PhaseChange { .. } => {}
205                AgentTraceEvent::LoopComplete {
206                    status: s,
207                    iterations: i,
208                    total_duration_ms: d,
209                    ..
210                } => {
211                    status = s.clone();
212                    iterations = *i;
213                    total_duration_ms = *d;
214                }
215                AgentTraceEvent::SchemaRetry { .. } => {}
216            }
217        }
218
219        serde_json::json!({
220            "status": status,
221            "iterations": iterations,
222            "total_duration_ms": total_duration_ms,
223            "llm_calls": llm_calls,
224            "tool_executions": tool_executions,
225            "tool_rejections": tool_rejections,
226            "interventions": interventions,
227            "compactions": compactions,
228            "total_input_tokens": total_input_tokens,
229            "total_output_tokens": total_output_tokens,
230            "total_llm_duration_ms": total_llm_duration_ms,
231            "total_tool_duration_ms": total_tool_duration_ms,
232            "tools_used": tools_used,
233        })
234    })
235}
236
237/// Reset agent trace state. Call between test runs.
238pub(crate) fn reset_agent_trace_state() {
239    AGENT_TRACE.with(|v| v.borrow_mut().clear());
240}