agtrace_engine/session/
extensions.rs

1//! Extension traits for session analysis.
2//!
3//! These traits provide complex computation logic for session types
4//! defined in agtrace-types, keeping the types crate lightweight and
5//! focused on data structures while engine handles analysis logic.
6
7use super::types::{AgentSession, TurnMetrics};
8
9/// Extension trait for `AgentSession` providing analysis and metrics computation.
10pub trait SessionAnalysisExt {
11    /// Compute presentation metrics for all turns
12    fn compute_turn_metrics(&self, max_context: Option<u32>) -> Vec<TurnMetrics>;
13}
14
15impl SessionAnalysisExt for AgentSession {
16    fn compute_turn_metrics(&self, max_context: Option<u32>) -> Vec<TurnMetrics> {
17        let mut cumulative_total = 0u32;
18        let mut metrics = Vec::new();
19
20        // Filter out turns that have no usage data yet (streaming just started)
21        // This prevents the context bar from showing delta=0 for in-progress turns
22        let turns_with_usage: Vec<_> = self
23            .turns
24            .iter()
25            .enumerate()
26            .filter(|(_, turn)| turn.steps.iter().any(|s| s.usage.is_some()))
27            .collect();
28
29        let total_valid_turns = turns_with_usage.len();
30
31        for (metric_idx, (original_idx, turn)) in turns_with_usage.into_iter().enumerate() {
32            let turn_end_cumulative = turn.cumulative_total_tokens(cumulative_total);
33            let delta = turn_end_cumulative.saturating_sub(cumulative_total);
34            let prev_total = cumulative_total;
35
36            // Last turn with usage is active during streaming
37            let is_active = metric_idx == total_valid_turns.saturating_sub(1);
38
39            metrics.push(TurnMetrics {
40                turn_index: original_idx,
41                prev_total,
42                delta,
43                is_heavy: TurnMetrics::is_delta_heavy(delta, max_context),
44                is_active,
45            });
46
47            cumulative_total = turn_end_cumulative;
48        }
49
50        metrics
51    }
52}