1use agtrace_types::{
2 MessagePayload, ReasoningPayload, TokenUsagePayload, ToolCallPayload, ToolResultPayload,
3 UserPayload,
4};
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct AgentSession {
14 pub session_id: Uuid,
15 pub start_time: DateTime<Utc>,
16 pub end_time: Option<DateTime<Utc>>,
17
18 pub turns: Vec<AgentTurn>,
19
20 pub stats: SessionStats,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AgentTurn {
28 pub id: Uuid,
29 pub timestamp: DateTime<Utc>,
30
31 pub user: UserMessage,
33
34 pub steps: Vec<AgentStep>,
37
38 pub stats: TurnStats,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct AgentStep {
46 pub id: Uuid,
47 pub timestamp: DateTime<Utc>,
48
49 pub reasoning: Option<ReasoningBlock>,
53
54 pub message: Option<MessageBlock>,
56
57 pub tools: Vec<ToolExecution>,
61
62 pub usage: Option<TokenUsagePayload>,
64 pub is_failed: bool,
65 pub status: StepStatus,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71pub enum StepStatus {
72 Done,
74 InProgress,
76 Failed,
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct ToolExecution {
87 pub call: ToolCallBlock,
88
89 pub result: Option<ToolResultBlock>,
91
92 pub duration_ms: Option<i64>,
94
95 pub is_error: bool,
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct UserMessage {
103 pub event_id: Uuid,
104 pub content: UserPayload,
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct ReasoningBlock {
109 pub event_id: Uuid,
110 pub content: ReasoningPayload,
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct MessageBlock {
115 pub event_id: Uuid,
116 pub content: MessagePayload,
117}
118
119#[derive(Debug, Clone, Serialize, Deserialize)]
120pub struct ToolCallBlock {
121 pub event_id: Uuid,
122 pub timestamp: DateTime<Utc>,
123 pub provider_call_id: Option<String>,
124 pub content: ToolCallPayload,
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct ToolResultBlock {
129 pub event_id: Uuid,
130 pub timestamp: DateTime<Utc>,
131 pub tool_call_id: Uuid,
132 pub content: ToolResultPayload,
133}
134
135#[derive(Debug, Clone, Serialize, Deserialize, Default)]
138pub struct SessionStats {
139 pub total_turns: usize,
140 pub duration_seconds: i64,
141 pub total_tokens: i64,
142}
143
144#[derive(Debug, Clone, Serialize, Deserialize, Default)]
145pub struct TurnStats {
146 pub duration_ms: i64,
147 pub step_count: usize,
148 pub total_tokens: i32,
149}
150
151#[derive(Debug, Clone)]
157pub struct TurnMetrics {
158 pub turn_index: usize,
159 pub prev_total: u32,
160 pub delta: u32,
161 pub is_heavy: bool,
162 pub is_active: bool,
163}
164
165impl TurnMetrics {
166 pub fn heavy_threshold(max_context: Option<u32>) -> u32 {
168 max_context.map(|mc| mc / 10).unwrap_or(15000)
169 }
170
171 pub fn is_delta_heavy(delta: u32, max_context: Option<u32>) -> bool {
173 delta >= Self::heavy_threshold(max_context)
174 }
175}
176
177impl AgentTurn {
178 pub fn cumulative_input_tokens(&self, fallback: u32) -> u32 {
181 self.steps
182 .iter()
183 .rev()
184 .find_map(|step| step.usage.as_ref())
185 .map(|usage| {
186 (usage.input_tokens
187 + usage
188 .details
189 .as_ref()
190 .and_then(|d| d.cache_creation_input_tokens)
191 .unwrap_or(0)
192 + usage
193 .details
194 .as_ref()
195 .and_then(|d| d.cache_read_input_tokens)
196 .unwrap_or(0)) as u32
197 })
198 .unwrap_or(fallback)
199 }
200
201 pub fn cumulative_total_tokens(&self, fallback: u32) -> u32 {
204 self.steps
205 .iter()
206 .rev()
207 .find_map(|step| step.usage.as_ref())
208 .map(|usage| {
209 (usage.input_tokens
210 + usage
211 .details
212 .as_ref()
213 .and_then(|d| d.cache_creation_input_tokens)
214 .unwrap_or(0)
215 + usage
216 .details
217 .as_ref()
218 .and_then(|d| d.cache_read_input_tokens)
219 .unwrap_or(0)
220 + usage.output_tokens) as u32
221 })
222 .unwrap_or(fallback)
223 }
224
225 pub fn is_active(&self) -> bool {
231 const LOOKBACK_STEPS: usize = 3;
232
233 self.steps
234 .iter()
235 .rev()
236 .take(LOOKBACK_STEPS)
237 .any(|step| matches!(step.status, StepStatus::InProgress))
238 }
239}
240
241impl AgentSession {
242 pub fn compute_turn_metrics(&self, max_context: Option<u32>) -> Vec<TurnMetrics> {
244 let mut cumulative_total = 0u32;
245 let mut metrics = Vec::new();
246 let total_turns = self.turns.len();
247
248 for (idx, turn) in self.turns.iter().enumerate() {
249 let turn_end_cumulative = turn.cumulative_total_tokens(cumulative_total);
250 let delta = turn_end_cumulative.saturating_sub(cumulative_total);
251 let prev_total = cumulative_total;
252
253 let is_active = if idx == total_turns.saturating_sub(1) {
256 true
259 } else {
260 false
261 };
262
263 metrics.push(TurnMetrics {
264 turn_index: idx,
265 prev_total,
266 delta,
267 is_heavy: TurnMetrics::is_delta_heavy(delta, max_context),
268 is_active,
269 });
270
271 cumulative_total = turn_end_cumulative;
272 }
273
274 metrics
275 }
276}