agtrace_engine/session/
types.rs

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// ==========================================
10// 1. Session (entire conversation)
11// ==========================================
12#[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// ==========================================
24// 2. Turn (user-initiated interaction unit)
25// ==========================================
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AgentTurn {
28    pub id: Uuid,
29    pub timestamp: DateTime<Utc>,
30
31    /// Turn trigger (Input)
32    pub user: UserMessage,
33
34    /// Agent autonomous operation cycle (Steps)
35    /// Single step for simple conversation, multiple steps for autonomous agents
36    pub steps: Vec<AgentStep>,
37
38    pub stats: TurnStats,
39}
40
41// ==========================================
42// 3. Step (single LLM inference + execution unit)
43// ==========================================
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct AgentStep {
46    pub id: Uuid,
47    pub timestamp: DateTime<Utc>,
48
49    // --- Phase 1: Generation (Agent Outputs) ---
50    // These are generated in parallel or in arbitrary order before seeing tool results
51    /// Reasoning (CoT)
52    pub reasoning: Option<ReasoningBlock>,
53
54    /// Text message (answer to user, or declaration of tool execution)
55    pub message: Option<MessageBlock>,
56
57    // --- Phase 2: Execution (System Outputs) ---
58    /// Tool execution pairs (Call + Result)
59    /// Calls are generated in Phase 1, but managed here as pairs with Results
60    pub tools: Vec<ToolExecution>,
61
62    // --- Meta ---
63    pub usage: Option<TokenUsagePayload>,
64    pub is_failed: bool,
65    pub status: StepStatus,
66}
67
68/// Step completion status
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71pub enum StepStatus {
72    /// Step completed successfully (has Message or all tools have results)
73    Done,
74    /// Step is waiting for tool results or next action
75    InProgress,
76    /// Step failed with errors
77    Failed,
78}
79
80// ==========================================
81// Components
82// ==========================================
83
84/// Single tool execution unit (Call -> Result)
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct ToolExecution {
87    pub call: ToolCallBlock,
88
89    /// Execution result (None if incomplete or lost)
90    pub result: Option<ToolResultBlock>,
91
92    /// Latency (result.timestamp - call.timestamp)
93    pub duration_ms: Option<i64>,
94
95    /// Whether this individual tool execution failed
96    pub is_error: bool,
97}
98
99// --- ID Wrappers ---
100
101#[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// --- Stats ---
136
137#[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// ==========================================
152// Computed Metrics (for presentation)
153// ==========================================
154
155/// Computed metrics for a turn, used for presentation layer
156#[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    /// Calculate heavy threshold: 10% of max context, or fallback to 15k tokens
167    pub fn heavy_threshold(max_context: Option<u32>) -> u32 {
168        max_context.map(|mc| mc / 10).unwrap_or(15000)
169    }
170
171    /// Check if a delta is considered heavy
172    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    /// Calculate cumulative input tokens at the end of this turn
179    /// Falls back to `fallback` if no usage data found
180    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    /// Check if this turn is currently active
202    ///
203    /// A turn is active if any of the recent steps are in progress.
204    /// Looking at multiple steps provides stability during step transitions
205    /// (e.g., when a step completes but the next one hasn't started yet).
206    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    /// Compute presentation metrics for all turns
219    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            // Last turn is always active during streaming to avoid flicker
230            // when steps transition between InProgress and Done
231            let is_active = if idx == total_turns.saturating_sub(1) {
232                // For streaming sessions, last turn is active regardless of step status
233                // This prevents "CURRENT TURN" display from flickering during step transitions
234                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}