1use std::cell::RefCell;
2
3#[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
21pub fn enable_tracing() {
23 LLM_TRACING_ENABLED.with(|v| *v.borrow_mut() = true);
24}
25
26pub fn take_trace() -> Vec<LlmTraceEntry> {
28 LLM_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
29}
30
31pub fn peek_trace() -> Vec<LlmTraceEntry> {
33 LLM_TRACE.with(|v| v.borrow().clone())
34}
35
36pub 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
53pub(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#[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 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
138pub(crate) fn emit_agent_event(event: AgentTraceEvent) {
140 AGENT_TRACE.with(|v| v.borrow_mut().push(event));
141}
142
143pub fn take_agent_trace() -> Vec<AgentTraceEvent> {
145 AGENT_TRACE.with(|v| std::mem::take(&mut *v.borrow_mut()))
146}
147
148pub fn peek_agent_trace() -> Vec<AgentTraceEvent> {
150 AGENT_TRACE.with(|v| v.borrow().clone())
151}
152
153pub 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
237pub(crate) fn reset_agent_trace_state() {
239 AGENT_TRACE.with(|v| v.borrow_mut().clear());
240}