Skip to main content

harn_vm/llm/
trace.rs

1use std::cell::RefCell;
2
3/// A single LLM call trace entry.
4#[derive(Debug, Clone)]
5pub struct LlmTraceEntry {
6    pub model: String,
7    pub input_tokens: i64,
8    pub output_tokens: i64,
9    pub duration_ms: u64,
10}
11
12thread_local! {
13    static LLM_TRACE: RefCell<Vec<LlmTraceEntry>> = const { RefCell::new(Vec::new()) };
14    static LLM_TRACING_ENABLED: RefCell<bool> = const { RefCell::new(false) };
15}
16
17/// Enable LLM tracing for the current thread.
18pub fn enable_tracing() {
19    LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = true);
20}
21
22/// Get and clear the trace log.
23pub fn take_trace() -> Vec<LlmTraceEntry> {
24    LLM_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
25}
26
27/// Clone the current trace log without consuming it.
28pub fn peek_trace() -> Vec<LlmTraceEntry> {
29    LLM_TRACE.with(|v| v.borrow().clone())
30}
31
32/// Summarize trace usage without consuming entries.
33pub fn peek_trace_summary() -> (i64, i64, i64, i64) {
34    LLM_TRACE.with(|v| {
35        let entries = v.borrow();
36        let mut input = 0i64;
37        let mut output = 0i64;
38        let mut duration = 0i64;
39        let count = entries.len() as i64;
40        for e in entries.iter() {
41            input += e.input_tokens;
42            output += e.output_tokens;
43            duration += e.duration_ms as i64;
44        }
45        (input, output, duration, count)
46    })
47}
48
49/// Reset thread-local trace state. Call between test runs.
50pub(crate) fn reset_trace_state() {
51    LLM_TRACE.with(|v| v.borrow_mut().clear());
52    LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = false);
53}
54
55pub(crate) fn trace_llm_call(entry: LlmTraceEntry) {
56    LLM_TRACING_ENABLED.with(|enabled| {
57        if *enabled.borrow() {
58            LLM_TRACE.with(|v| v.borrow_mut().push(entry));
59        }
60    });
61}
62
63/// Fine-grained event emitted during agent loop execution. Captures tool
64/// calls, LLM calls, interventions, compaction, and phase changes so
65/// downstream consumers (portal, IDE hosts, cloud runners) can display
66/// execution traces without reconstructing them from raw JSON.
67#[derive(Debug, Clone, serde::Serialize)]
68#[serde(tag = "type", rename_all = "snake_case")]
69pub enum AgentTraceEvent {
70    LlmCall {
71        call_id: String,
72        model: String,
73        input_tokens: i64,
74        output_tokens: i64,
75        cache_tokens: i64,
76        duration_ms: u64,
77        iteration: usize,
78    },
79    ToolExecution {
80        tool_name: String,
81        tool_use_id: String,
82        duration_ms: u64,
83        status: String,
84        classification: String,
85        iteration: usize,
86    },
87    ToolRejected {
88        tool_name: String,
89        reason: String,
90        iteration: usize,
91    },
92    LoopIntervention {
93        tool_name: String,
94        kind: String,
95        count: usize,
96        iteration: usize,
97    },
98    ContextCompaction {
99        archived_messages: usize,
100        new_summary_len: usize,
101        iteration: usize,
102    },
103    PhaseChange {
104        from_phase: String,
105        to_phase: String,
106        iteration: usize,
107    },
108    LoopComplete {
109        status: String,
110        iterations: usize,
111        total_duration_ms: u64,
112        tools_used: Vec<String>,
113        successful_tools: Vec<String>,
114    },
115    /// Emitted when `llm_call` re-prompts the model after the previous
116    /// response failed `output_schema` validation. One event per retry;
117    /// `attempt` counts retries (the initial call is attempt 0 and
118    /// produces no event; the first retry emits `attempt: 1`).
119    ///
120    /// The retry does **not** persist the invalid response — the
121    /// original messages are replayed with a single appended user-role
122    /// correction that cites the validation errors and schema. That
123    /// correction text is surfaced here as `correction_prompt` so
124    /// transcripts show both why the retry happened and what was sent.
125    SchemaRetry {
126        attempt: usize,
127        errors: Vec<String>,
128        nudge_used: bool,
129        correction_prompt: String,
130    },
131    /// Emitted when `llm_call` aborts a streaming provider response
132    /// because the partial JSON content can no longer satisfy
133    /// `output_schema`. `chunks_consumed` counts text-delta chunks seen
134    /// before the abort; `provider` / `model` track the route that fired
135    /// so cost dashboards can attribute the savings.
136    SchemaStreamAborted {
137        provider: String,
138        model: String,
139        reason: String,
140        path: String,
141        chunks_consumed: usize,
142    },
143    TypedCheckpoint {
144        name: String,
145        status: String,
146        checkpoint_attempts: usize,
147        llm_attempts: usize,
148        error_category: Option<String>,
149        errors: Vec<String>,
150        repaired: bool,
151        final_accepted: bool,
152        raw_text: String,
153    },
154    NativeToolFallback {
155        iteration: usize,
156        accepted: bool,
157        policy: String,
158        fallback_index: usize,
159        tool_call_count: usize,
160    },
161    EmptyCompletionRetry {
162        iteration: usize,
163        attempt: usize,
164        provider: String,
165        model: String,
166        reason: String,
167        duration_ms: u64,
168        error: String,
169    },
170    /// Emitted when a `models:`/`ladder:` model ladder advances from one rung
171    /// to the next because the current rung hit a transport-class failure
172    /// (connection/timeout/429/5xx/circuit_open). Schema-validation failures
173    /// never emit this — they re-ask the SAME rung's model. `from_index` is
174    /// the 0-based ladder position that failed; `category` is the failover
175    /// error category that drove the advance.
176    ModelsAdvance {
177        from_index: usize,
178        from_model: String,
179        to_model: String,
180        category: String,
181    },
182}
183
184thread_local! {
185    static AGENT_TRACE: RefCell<Vec<AgentTraceEvent>> = const { RefCell::new(Vec::new()) };
186}
187
188/// Emit an agent trace event.
189pub(crate) fn emit_agent_event(event: AgentTraceEvent) {
190    AGENT_TRACE.with(|v| v.borrow_mut().push(event));
191}
192
193/// Get and clear the agent trace log.
194pub fn take_agent_trace() -> Vec<AgentTraceEvent> {
195    AGENT_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
196}
197
198/// Clone the current agent trace log without consuming it.
199pub fn peek_agent_trace() -> Vec<AgentTraceEvent> {
200    AGENT_TRACE.with(|v| v.borrow().clone())
201}
202
203/// Produce a rolled-up summary of agent trace events as JSON.
204pub fn agent_trace_summary() -> serde_json::Value {
205    AGENT_TRACE.with(|v| {
206        let events = v.borrow();
207        let mut llm_calls = 0usize;
208        let mut tool_executions = 0usize;
209        let mut tool_rejections = 0usize;
210        let mut interventions = 0usize;
211        let mut compactions = 0usize;
212        let mut native_text_tool_fallbacks = 0usize;
213        let mut native_text_tool_fallback_rejections = 0usize;
214        let mut empty_completion_retries = 0usize;
215        let mut models_advances = 0usize;
216        let mut schema_stream_aborts = 0usize;
217        let mut typed_checkpoints = 0usize;
218        let mut typed_checkpoint_failures = 0usize;
219        let mut total_input_tokens = 0i64;
220        let mut total_output_tokens = 0i64;
221        let mut total_llm_duration_ms = 0u64;
222        let mut total_tool_duration_ms = 0u64;
223        let mut tools_used: Vec<String> = Vec::new();
224        let mut status = "unknown".to_string();
225        let mut iterations = 0usize;
226        let mut total_duration_ms = 0u64;
227
228        for event in events.iter() {
229            match event {
230                AgentTraceEvent::LlmCall {
231                    input_tokens,
232                    output_tokens,
233                    duration_ms,
234                    ..
235                } => {
236                    llm_calls += 1;
237                    total_input_tokens += input_tokens;
238                    total_output_tokens += output_tokens;
239                    total_llm_duration_ms += duration_ms;
240                }
241                AgentTraceEvent::ToolExecution {
242                    tool_name,
243                    duration_ms,
244                    ..
245                } => {
246                    tool_executions += 1;
247                    total_tool_duration_ms += duration_ms;
248                    if !tools_used.contains(tool_name) {
249                        tools_used.push(tool_name.clone());
250                    }
251                }
252                AgentTraceEvent::ToolRejected { .. } => {
253                    tool_rejections += 1;
254                }
255                AgentTraceEvent::LoopIntervention { .. } => {
256                    interventions += 1;
257                }
258                AgentTraceEvent::ContextCompaction { .. } => {
259                    compactions += 1;
260                }
261                AgentTraceEvent::PhaseChange { .. } => {}
262                AgentTraceEvent::LoopComplete {
263                    status: s,
264                    iterations: i,
265                    total_duration_ms: d,
266                    ..
267                } => {
268                    status = s.clone();
269                    iterations = *i;
270                    total_duration_ms = *d;
271                }
272                AgentTraceEvent::SchemaRetry { .. } => {}
273                AgentTraceEvent::SchemaStreamAborted { .. } => {
274                    schema_stream_aborts += 1;
275                }
276                AgentTraceEvent::TypedCheckpoint { final_accepted, .. } => {
277                    typed_checkpoints += 1;
278                    if !final_accepted {
279                        typed_checkpoint_failures += 1;
280                    }
281                }
282                AgentTraceEvent::NativeToolFallback { accepted, .. } => {
283                    native_text_tool_fallbacks += 1;
284                    if !accepted {
285                        native_text_tool_fallback_rejections += 1;
286                    }
287                }
288                AgentTraceEvent::EmptyCompletionRetry { .. } => {
289                    empty_completion_retries += 1;
290                }
291                AgentTraceEvent::ModelsAdvance { .. } => {
292                    models_advances += 1;
293                }
294            }
295        }
296
297        serde_json::json!({
298            "status": status,
299            "iterations": iterations,
300            "total_duration_ms": total_duration_ms,
301            "llm_calls": llm_calls,
302            "tool_executions": tool_executions,
303            "tool_rejections": tool_rejections,
304            "interventions": interventions,
305            "compactions": compactions,
306            "native_text_tool_fallbacks": native_text_tool_fallbacks,
307            "native_text_tool_fallback_rejections": native_text_tool_fallback_rejections,
308            "empty_completion_retries": empty_completion_retries,
309            "models_advances": models_advances,
310            "schema_stream_aborts": schema_stream_aborts,
311            "typed_checkpoints": typed_checkpoints,
312            "typed_checkpoint_failures": typed_checkpoint_failures,
313            "total_input_tokens": total_input_tokens,
314            "total_output_tokens": total_output_tokens,
315            "total_llm_duration_ms": total_llm_duration_ms,
316            "total_tool_duration_ms": total_tool_duration_ms,
317            "tools_used": tools_used,
318        })
319    })
320}
321
322/// Reset agent trace state. Call between test runs.
323pub(crate) fn reset_agent_trace_state() {
324    AGENT_TRACE.with(|v| v.borrow_mut().clear());
325}