yapi_core/agent_session/
stats.rs

1//! Context use, token and cost totals, and session statistics.
2
3use yapi_types::message::{ContentBlock, Message, StopReason};
4use yapi_types::session::FileEntry;
5
6use crate::compaction::{calculate_context_tokens, estimate_projected_context_tokens};
7use crate::session::build_projection;
8
9use super::AgentSession;
10
11/// How full the context window is.
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub struct ContextUsage {
14    /// Estimated context tokens, unknown right after a compaction.
15    pub tokens: Option<u64>,
16    /// The model's context window.
17    pub context_window: u64,
18}
19
20impl ContextUsage {
21    /// Percent of the window in use, when known.
22    pub fn percent(&self) -> Option<f64> {
23        self.tokens
24            .map(|tokens| tokens as f64 / self.context_window as f64 * 100.0)
25    }
26}
27
28/// Session-wide token and cost totals.
29#[derive(Clone, Copy, Debug, Default, PartialEq)]
30pub struct UsageTotals {
31    /// Input tokens.
32    pub input: u64,
33    /// Output tokens.
34    pub output: u64,
35    /// Cache read tokens.
36    pub cache_read: u64,
37    /// Cache write tokens.
38    pub cache_write: u64,
39    /// Cost in US dollars.
40    pub cost: f64,
41    /// Cache hit rate of the latest assistant message, in percent.
42    pub cache_hit_rate: Option<f64>,
43}
44
45/// Message counts and per-model cost, as `/session` shows them.
46#[derive(Clone, Debug, Default, PartialEq)]
47pub struct SessionStats {
48    /// Message entries.
49    pub total_messages: usize,
50    /// User messages.
51    pub user_messages: usize,
52    /// Assistant messages.
53    pub assistant_messages: usize,
54    /// Tool calls in assistant messages.
55    pub tool_calls: usize,
56    /// Tool results.
57    pub tool_results: usize,
58    /// Cost and tokens by `provider/model`, costliest first. Summaries and
59    /// tool usage are grouped as `Tools/summaries`.
60    pub breakdown: Vec<(String, f64, u64)>,
61}
62
63impl AgentSession {
64    /// Context use of the current branch: `tokens` is `None` after a compaction
65    /// until the model responds again. `None` without a model or window.
66    pub fn context_usage(&self) -> Option<ContextUsage> {
67        let model = self.model()?;
68        let context_window = model.context_window;
69        if context_window == 0 {
70            return None;
71        }
72        self.with_session(|session| {
73            let branch = session.branch_path(None);
74            let projection = build_projection(&branch);
75            let latest_compaction = branch
76                .iter()
77                .rposition(|entry| matches!(entry, FileEntry::Compaction(_)));
78            if let Some(compaction) = latest_compaction {
79                let has_usage = |entry: &&FileEntry| {
80                    let Some(id) = entry.meta().map(|meta| meta.id.as_str()) else {
81                        return false;
82                    };
83                    projection.entries.iter().any(|projected| {
84                        projected.source.meta().is_some_and(|meta| meta.id == id)
85                            && projected.messages.iter().any(|message| {
86                                matches!(message, Message::Assistant(assistant)
87                                    if !matches!(assistant.stop_reason, StopReason::Aborted | StopReason::Error)
88                                        && calculate_context_tokens(&assistant.usage) > 0)
89                            })
90                    })
91                };
92                if !branch[compaction + 1..].iter().any(has_usage) {
93                    return Some(ContextUsage {
94                        tokens: None,
95                        context_window,
96                    });
97                }
98            }
99            let tokens = estimate_projected_context_tokens(&projection, &branch).tokens;
100            Some(ContextUsage {
101                tokens: Some(tokens),
102                context_window,
103            })
104        })
105    }
106
107    /// Token and cost totals over every entry of the session file, and the
108    /// cache hit rate of the latest assistant message.
109    pub fn usage_totals(&self) -> UsageTotals {
110        self.with_session(|session| {
111            let mut totals = UsageTotals::default();
112            for entry in session.entries() {
113                let usage = match entry {
114                    FileEntry::Usage(entry) => Some(&entry.usage),
115                    FileEntry::Message(entry) => match &entry.message {
116                        Message::Assistant(assistant) => {
117                            let usage = &assistant.usage;
118                            let prompt = usage.input + usage.cache_read + usage.cache_write;
119                            totals.cache_hit_rate = (prompt > 0)
120                                .then(|| usage.cache_read as f64 / prompt as f64 * 100.0);
121                            Some(usage)
122                        }
123                        Message::ToolResult(result) => result.usage.as_ref(),
124                        _ => None,
125                    },
126                    FileEntry::Compaction(entry) => entry.usage.as_ref(),
127                    FileEntry::BranchSummary(entry) => entry.usage.as_ref(),
128                    _ => None,
129                };
130                if let Some(usage) = usage {
131                    totals.input += usage.input;
132                    totals.output += usage.output;
133                    totals.cache_read += usage.cache_read;
134                    totals.cache_write += usage.cache_write;
135                    totals.cost += usage.cost.total;
136                }
137            }
138            totals
139        })
140    }
141
142    /// pi's `getSessionStats` counts and `getUsageCostBreakdown`, over every
143    /// entry of the session file.
144    pub fn session_stats(&self) -> SessionStats {
145        self.with_session(|session| {
146            let mut stats = SessionStats::default();
147            let mut breakdown: Vec<(String, f64, u64)> = Vec::new();
148            let mut add = |key: String, usage: &yapi_types::message::Usage| {
149                let tokens = usage.input + usage.output + usage.cache_read + usage.cache_write;
150                match breakdown
151                    .iter_mut()
152                    .find(|(existing, _, _)| *existing == key)
153                {
154                    Some((_, cost, total)) => {
155                        *cost += usage.cost.total;
156                        *total += tokens;
157                    }
158                    None => breakdown.push((key, usage.cost.total, tokens)),
159                }
160            };
161            for entry in session.entries() {
162                match entry {
163                    FileEntry::Usage(entry) => {
164                        add(format!("{}/{}", entry.provider, entry.model), &entry.usage);
165                    }
166                    FileEntry::Compaction(entry) => {
167                        if let Some(usage) = &entry.usage {
168                            add("Tools/summaries".into(), usage);
169                        }
170                    }
171                    FileEntry::BranchSummary(entry) => {
172                        if let Some(usage) = &entry.usage {
173                            add("Tools/summaries".into(), usage);
174                        }
175                    }
176                    FileEntry::Message(entry) => {
177                        stats.total_messages += 1;
178                        match &entry.message {
179                            Message::User(_) => stats.user_messages += 1,
180                            Message::ToolResult(result) => {
181                                stats.tool_results += 1;
182                                if let Some(usage) = &result.usage {
183                                    add("Tools/summaries".into(), usage);
184                                }
185                            }
186                            Message::Assistant(assistant) => {
187                                stats.assistant_messages += 1;
188                                stats.tool_calls += assistant
189                                    .content
190                                    .iter()
191                                    .filter(|block| matches!(block, ContentBlock::ToolCall(_)))
192                                    .count();
193                                let model = assistant
194                                    .response_model
195                                    .as_deref()
196                                    .unwrap_or(&assistant.model);
197                                add(format!("{}/{model}", assistant.provider), &assistant.usage);
198                            }
199                            _ => {}
200                        }
201                    }
202                    _ => {}
203                }
204            }
205            breakdown.retain(|(_, cost, tokens)| *cost > 0.0 || *tokens > 0);
206            breakdown.sort_by(|a, b| b.1.total_cmp(&a.1));
207            stats.breakdown = breakdown;
208            stats
209        })
210    }
211}
yapi_core/agent_session/stats.rs

yapi_core/agent_session/
stats.rs
