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 is_active(&self) -> bool {
207 const LOOKBACK_STEPS: usize = 3;
208
209 self.steps
210 .iter()
211 .rev()
212 .take(LOOKBACK_STEPS)
213 .any(|step| matches!(step.status, StepStatus::InProgress))
214 }
215}
216
217impl AgentSession {
218 pub fn compute_turn_metrics(&self, max_context: Option<u32>) -> Vec<TurnMetrics> {
220 let mut cumulative_input = 0u32;
221 let mut metrics = Vec::new();
222 let total_turns = self.turns.len();
223
224 for (idx, turn) in self.turns.iter().enumerate() {
225 let turn_end_cumulative = turn.cumulative_input_tokens(cumulative_input);
226 let delta = turn_end_cumulative.saturating_sub(cumulative_input);
227 let prev_total = cumulative_input;
228
229 let is_active = if idx == total_turns.saturating_sub(1) {
232 true
235 } else {
236 false
237 };
238
239 metrics.push(TurnMetrics {
240 turn_index: idx,
241 prev_total,
242 delta,
243 is_heavy: TurnMetrics::is_delta_heavy(delta, max_context),
244 is_active,
245 });
246
247 cumulative_input = turn_end_cumulative;
248 }
249
250 metrics
251 }
252}