1use std::cell::RefCell;
2
3#[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
17pub fn enable_tracing() {
19 LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = true);
20}
21
22pub fn take_trace() -> Vec<LlmTraceEntry> {
24 LLM_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
25}
26
27pub fn peek_trace() -> Vec<LlmTraceEntry> {
29 LLM_TRACE.with(|v| v.borrow().clone())
30}
31
32pub 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
49pub(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#[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 SchemaRetry {
126 attempt: usize,
127 errors: Vec<String>,
128 nudge_used: bool,
129 correction_prompt: String,
130 },
131 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 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
188pub(crate) fn emit_agent_event(event: AgentTraceEvent) {
190 AGENT_TRACE.with(|v| v.borrow_mut().push(event));
191}
192
193pub fn take_agent_trace() -> Vec<AgentTraceEvent> {
195 AGENT_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
196}
197
198pub fn peek_agent_trace() -> Vec<AgentTraceEvent> {
200 AGENT_TRACE.with(|v| v.borrow().clone())
201}
202
203pub 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
322pub(crate) fn reset_agent_trace_state() {
324 AGENT_TRACE.with(|v| v.borrow_mut().clear());
325}