pi/core/agent_session/
stats.rs

1//! Session statistics + context usage impls.
2//!
3//! Ports `.references/pi/packages/coding-agent/src/core/agent-session.ts`
4//! `getSessionStats`, `getContextUsage`, and the
5//! `ContextUsage` / `SessionStats` types.
6//!
7//! Behaviour preserved from the TypeScript contract:
8//! - `get_session_stats` aggregates over **all** session entries (including
9//!   history that was compacted away), so token / cost totals reflect what
10//!   was actually billed across the session.
11//! - `get_context_usage` reports `tokens: null` / `percent: null` when the
12//!   latest compaction has no post-compaction assistant usage yet (the next
13//!   LLM response must establish the new baseline).
14//! - Cost / token sums use the canonical [`pi_ai::Usage`] fields, with the
15//!   total falling back to `input + output + cache_read + cache_write` when
16//!   `total_tokens` is zero (TypeScript `calculateContextTokens`).
17//!
18//! Lock order: only `AgentSessionInner` (briefly, to read mirrors). The
19//! session manager async mutex is acquired for entry enumeration.
20
21use pi_ai::{AssistantContent, Message, Model, StopReason};
22
23use crate::core::compaction::{calculate_context_tokens, estimate_context_tokens};
24use crate::core::sessions::{SessionEntry, get_latest_compaction_entry};
25
26use super::AgentSession;
27
28/// Aggregate session statistics (TypeScript `SessionStats`).
29#[derive(Clone, Debug, PartialEq)]
30pub struct SessionStats {
31    /// Session file path, if any.
32    pub session_file: Option<String>,
33    /// Session id.
34    pub session_id: String,
35    /// Number of user messages.
36    pub user_messages: u64,
37    /// Number of assistant messages.
38    pub assistant_messages: u64,
39    /// Number of tool calls across all assistant messages.
40    pub tool_calls: u64,
41    /// Number of tool-result messages.
42    pub tool_results: u64,
43    /// Total messages (assistant + user + toolResult).
44    pub total_messages: u64,
45    /// Token totals.
46    pub tokens: SessionTokenTotals,
47    /// Total billed cost in US dollars.
48    pub cost: f64,
49    /// Context-usage snapshot, when computable.
50    pub context_usage: Option<ContextUsage>,
51}
52
53/// Token totals (TypeScript `SessionStats['tokens']`).
54#[derive(Clone, Copy, Debug, Default, PartialEq)]
55pub struct SessionTokenTotals {
56    /// Sum of assistant `usage.input`.
57    pub input: u64,
58    /// Sum of assistant `usage.output`.
59    pub output: u64,
60    /// Sum of assistant `usage.cacheRead`.
61    pub cache_read: u64,
62    /// Sum of assistant `usage.cacheWrite`.
63    pub cache_write: u64,
64    /// `input + output + cache_read + cache_write`.
65    pub total: u64,
66}
67
68/// Context-usage snapshot (TypeScript `ContextUsage`).
69///
70/// `tokens` / `percent` are `None` after a compaction until the next
71/// assistant response establishes a fresh usage baseline.
72#[derive(Clone, Copy, Debug, PartialEq)]
73pub struct ContextUsage {
74    /// Estimated context tokens (`None` when unknown after compaction).
75    pub tokens: Option<u64>,
76    /// Model context-window size.
77    pub context_window: u64,
78    /// `tokens / context_window * 100.0` (`None` when `tokens` is unknown).
79    pub percent: Option<f64>,
80}
81
82impl AgentSession {
83    /// Aggregate session statistics across all persisted entries.
84    ///
85    /// Counts / totals include compacted-away history so cost reflects what
86    /// was actually billed. See [`Self::get_context_usage`] for the live
87    /// context estimate used by the UI.
88    pub async fn get_session_stats(&self) -> SessionStats {
89        let (session_file, session_id, entries): (Option<String>, String, Vec<SessionEntry>) = {
90            let manager = self.session_manager.lock().await;
91            (
92                manager.get_session_file().map(str::to_owned),
93                manager.get_session_id().to_owned(),
94                manager.get_entries().into_iter().cloned().collect(),
95            )
96        };
97
98        let mut user_messages = 0u64;
99        let mut assistant_messages = 0u64;
100        let mut tool_results = 0u64;
101        let mut total_messages = 0u64;
102        let mut tool_calls = 0u64;
103        let mut input = 0u64;
104        let mut output = 0u64;
105        let mut cache_read = 0u64;
106        let mut cache_write = 0u64;
107        let mut cost = 0f64;
108
109        for entry in &entries {
110            let SessionEntry::Message(message_entry) = entry else {
111                continue;
112            };
113            total_messages = total_messages.saturating_add(1);
114            let message = &message_entry.message;
115            match message.as_llm() {
116                Some(Message::User(_)) => {
117                    user_messages = user_messages.saturating_add(1);
118                }
119                Some(Message::ToolResult(_)) => {
120                    tool_results = tool_results.saturating_add(1);
121                }
122                Some(Message::Assistant(assistant)) => {
123                    assistant_messages = assistant_messages.saturating_add(1);
124                    tool_calls = tool_calls.saturating_add(
125                        assistant
126                            .content
127                            .iter()
128                            .filter(|content| matches!(content, AssistantContent::ToolCall(_)))
129                            .count() as u64,
130                    );
131                    input = input.saturating_add(assistant.usage.input);
132                    output = output.saturating_add(assistant.usage.output);
133                    cache_read = cache_read.saturating_add(assistant.usage.cache_read);
134                    cache_write = cache_write.saturating_add(assistant.usage.cache_write);
135                    cost += assistant.usage.cost.total;
136                }
137                None => {}
138            }
139        }
140
141        let total = input
142            .saturating_add(output)
143            .saturating_add(cache_read)
144            .saturating_add(cache_write);
145
146        SessionStats {
147            session_file,
148            session_id,
149            user_messages,
150            assistant_messages,
151            tool_calls,
152            tool_results,
153            total_messages,
154            tokens: SessionTokenTotals {
155                input,
156                output,
157                cache_read,
158                cache_write,
159                total,
160            },
161            cost,
162            context_usage: self.get_context_usage().await,
163        }
164    }
165
166    /// Estimate current context usage against the active model's window.
167    ///
168    /// Returns `None` when the current model has no `contextWindow` (or no
169    /// model is set). After the latest compaction on the active branch, the
170    /// estimate is `None` until a post-compaction assistant response provides
171    /// a fresh usage baseline (TypeScript `getContextUsage`).
172    pub async fn get_context_usage(&self) -> Option<ContextUsage> {
173        let model = self.model();
174        let context_window = model.context_window;
175        if context_window == 0 {
176            return None;
177        }
178
179        // Branch + compaction boundary come from the persisted tree.
180        let branch: Vec<SessionEntry> = {
181            let manager = self.session_manager.lock().await;
182            manager.get_branch(None).into_iter().cloned().collect()
183        };
184
185        let branch_refs: Vec<&SessionEntry> = branch.iter().collect();
186        let latest_compaction = get_latest_compaction_entry(&branch_refs);
187
188        if let Some(compaction) = latest_compaction {
189            // `compaction` is `&CompactionEntry` borrowed from `branch_refs`.
190            // Find the same entry by pointer identity in the owned `branch`
191            // vector (the `&SessionEntry` pattern binds `c: &CompactionEntry`).
192            let compaction_index = branch
193                .iter()
194                .position(|entry| {
195                    matches!(
196                        entry,
197                        SessionEntry::Compaction(c) if std::ptr::eq(c, compaction)
198                    )
199                })
200                .unwrap_or(0);
201            let has_post_compaction_usage = branch.iter().enumerate().any(|(idx, entry)| {
202                if idx <= compaction_index {
203                    return false;
204                }
205                post_compaction_usage_tokens(entry).is_some()
206            });
207            if !has_post_compaction_usage {
208                return Some(ContextUsage {
209                    tokens: None,
210                    context_window,
211                    percent: None,
212                });
213            }
214        }
215
216        let messages = self.messages();
217        let estimate = estimate_context_tokens(&messages);
218        let tokens = estimate.tokens;
219        let percent = u64_as_f64(tokens) / u64_as_f64(context_window) * 100.0;
220        Some(ContextUsage {
221            tokens: Some(tokens),
222            context_window,
223            percent: Some(percent),
224        })
225    }
226}
227
228/// Convert a `u64` to `f64` without a precision-loss cast.
229///
230/// Splitting at the 32-bit boundary produces the same rounded binary value as
231/// Rust's primitive conversion while keeping both integer-to-float conversions
232/// lossless.
233fn u64_as_f64(value: u64) -> f64 {
234    let bytes = value.to_be_bytes();
235    let high = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
236    let low = u32::from_be_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
237    f64::from(high).mul_add(4_294_967_296.0, f64::from(low))
238}
239
240/// Tokens reported by an assistant entry usable as a post-compaction baseline.
241///
242/// Skips `aborted` / `error` stop reasons (TypeScript
243/// `assistant.stopReason !== "aborted" && assistant.stopReason !== "error"`)
244/// and zero-token usage records.
245fn post_compaction_usage_tokens(entry: &SessionEntry) -> Option<u64> {
246    let SessionEntry::Message(message_entry) = entry else {
247        return None;
248    };
249    let message = &message_entry.message;
250    let Message::Assistant(assistant) = message.as_llm()? else {
251        return None;
252    };
253    if matches!(
254        assistant.stop_reason,
255        StopReason::Aborted | StopReason::Error
256    ) {
257        return None;
258    }
259    let tokens = calculate_context_tokens(&assistant.usage);
260    if tokens == 0 {
261        return None;
262    }
263    Some(tokens)
264}
265
266/// Helper retained for sibling modules / future slices that need to read the
267/// active model's context window without downcasting the runtime.
268#[allow(dead_code)]
269fn model_context_window(model: &Model) -> u64 {
270    model.context_window
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn context_usage_tokens_when_window_zero_is_none() {
279        // Direct constructor test: a zero context window means we cannot
280        // estimate; the public getter returns `None`.
281        let usage = ContextUsage {
282            tokens: Some(100),
283            context_window: 0,
284            percent: None,
285        };
286        assert_eq!(usage.context_window, 0);
287    }
288
289    #[test]
290    fn token_totals_saturate_instead_of_overflow() {
291        let totals = SessionTokenTotals {
292            input: u64::MAX,
293            output: 1,
294            cache_read: 0,
295            cache_write: 0,
296            total: 0,
297        };
298        // The aggregation path uses saturating_add; just sanity-check that
299        // the helper struct accepts extreme values without panic.
300        assert_eq!(totals.input, u64::MAX);
301    }
302}
pi/core/agent_session/stats.rs

pi/core/agent_session/
stats.rs
