Skip to main content

oino_core/session/
stats.rs

1//! Session statistics, context usage and export (`getSessionStats`, `getContextUsage`,
2//! `exportToJsonl`, `exportToHtml`).
3
4use std::path::PathBuf;
5
6use oino_ai::{AgentMessage, AssistantContent, StopReason, Usage};
7use oino_session::{EntryKind, FileLine, SessionEntry};
8use serde::{Deserialize, Serialize};
9
10use super::{AgentSession, SessionOpError, lock};
11use crate::compaction::{calculate_context_tokens, estimate_projected_context_tokens};
12
13/// The context usage record.
14#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16pub struct ContextUsage {
17    /// Estimated context tokens (`null` right after compaction).
18    pub tokens: Option<u64>,
19    /// Context window.
20    pub context_window: u64,
21    /// Percent of the window (`null` with unknown tokens).
22    pub percent: Option<f64>,
23}
24
25/// Token totals of [`SessionStats`].
26#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(rename_all = "camelCase")]
28pub struct TokenTotals {
29    /// Input tokens.
30    pub input: u64,
31    /// Output tokens.
32    pub output: u64,
33    /// Cache read tokens.
34    pub cache_read: u64,
35    /// Cache write tokens.
36    pub cache_write: u64,
37    /// Sum of the four.
38    pub total: u64,
39}
40
41/// The session stats record.
42#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
43#[serde(rename_all = "camelCase")]
44pub struct SessionStats {
45    /// Session file.
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub session_file: Option<String>,
48    /// Session id.
49    pub session_id: String,
50    /// User messages.
51    pub user_messages: u64,
52    /// Assistant messages.
53    pub assistant_messages: u64,
54    /// Tool calls.
55    pub tool_calls: u64,
56    /// Tool results.
57    pub tool_results: u64,
58    /// Message entries.
59    pub total_messages: u64,
60    /// Token totals.
61    pub tokens: TokenTotals,
62    /// Total cost (USD).
63    pub cost: f64,
64    /// Context usage.
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub context_usage: Option<ContextUsage>,
67}
68
69fn add(totals: &mut TokenTotals, cost: &mut f64, usage: &Usage) {
70    totals.input += usage.input;
71    totals.output += usage.output;
72    totals.cache_read += usage.cache_read;
73    totals.cache_write += usage.cache_write;
74    *cost += usage.cost.total;
75}
76
77impl AgentSession {
78    /// Session stats (all entries, including compacted history and `usage` entries).
79    #[must_use]
80    pub fn session_stats(&self) -> SessionStats {
81        let mut stats = SessionStats {
82            session_file: self.session_file().map(|p| p.display().to_string()),
83            session_id: self.session_id(),
84            user_messages: 0,
85            assistant_messages: 0,
86            tool_calls: 0,
87            tool_results: 0,
88            total_messages: 0,
89            tokens: TokenTotals::default(),
90            cost: 0.0,
91            context_usage: None,
92        };
93        let entries = self.entries();
94        for entry in &entries {
95            match &entry.kind {
96                EntryKind::Usage { usage, .. } => add(&mut stats.tokens, &mut stats.cost, usage),
97                EntryKind::BranchSummary(summary) => {
98                    if let Some(usage) = &summary.usage {
99                        add(&mut stats.tokens, &mut stats.cost, usage);
100                    }
101                }
102                EntryKind::Compaction(compaction) => {
103                    if let Some(usage) = &compaction.usage {
104                        add(&mut stats.tokens, &mut stats.cost, usage);
105                    }
106                }
107                EntryKind::Message { message } => {
108                    stats.total_messages += 1;
109                    match message {
110                        AgentMessage::User(_) => stats.user_messages += 1,
111                        AgentMessage::ToolResult(result) => {
112                            stats.tool_results += 1;
113                            if let Some(usage) = &result.usage {
114                                add(&mut stats.tokens, &mut stats.cost, usage);
115                            }
116                        }
117                        AgentMessage::Assistant(assistant) => {
118                            stats.assistant_messages += 1;
119                            stats.tool_calls += assistant
120                                .content
121                                .iter()
122                                .filter(|b| matches!(b, AssistantContent::ToolCall(_)))
123                                .count() as u64;
124                            add(&mut stats.tokens, &mut stats.cost, &assistant.usage);
125                        }
126                        _ => {}
127                    }
128                }
129                _ => {}
130            }
131        }
132        stats.tokens.total = stats.tokens.input
133            + stats.tokens.output
134            + stats.tokens.cache_read
135            + stats.tokens.cache_write;
136        stats.context_usage = self.context_usage();
137        stats
138    }
139
140    /// Returns the context usage.
141    #[must_use]
142    pub fn context_usage(&self) -> Option<ContextUsage> {
143        let model = self.routed_model().or_else(|| self.model())?;
144        let context_window = model.context_window;
145        if context_window == 0 {
146            return None;
147        }
148        let manager = lock(&self.inner.manager);
149        let projection = manager.build_session_projection();
150        let branch = manager.branch(None);
151        if let Some(compaction_index) = branch
152            .iter()
153            .rposition(|e| matches!(e.kind, EntryKind::Compaction(_)))
154        {
155            let projected: Vec<&oino_session::EntryId> = projection
156                .entries
157                .iter()
158                .filter(|entry| {
159                    entry.messages.iter().any(|m| {
160                        matches!(m, AgentMessage::Assistant(a) if !matches!(a.stop_reason, StopReason::Aborted | StopReason::Error) && calculate_context_tokens(&a.usage) > 0)
161                    })
162                })
163                .map(|entry| &entry.source_entry.id)
164                .collect();
165            let has_post = branch[compaction_index + 1..]
166                .iter()
167                .any(|e| projected.contains(&&e.id));
168            if !has_post {
169                return Some(ContextUsage {
170                    tokens: None,
171                    context_window,
172                    percent: None,
173                });
174            }
175        }
176        let estimate = estimate_projected_context_tokens(&projection, &branch);
177        #[allow(clippy::cast_precision_loss, reason = "percent is approximate")]
178        let percent = estimate.tokens as f64 / context_window as f64 * 100.0;
179        Some(ContextUsage {
180            tokens: Some(estimate.tokens),
181            context_window,
182            percent: Some(percent),
183        })
184    }
185
186    /// A fresh header plus the current branch, re-linked linearly.
187    ///
188    /// # Errors
189    ///
190    /// Serialisation or write failures.
191    pub async fn export_jsonl(&self, path: Option<PathBuf>) -> Result<PathBuf, SessionOpError> {
192        let now = oino_session::IsoTimestamp::now();
193        let lines = {
194            let manager = lock(&self.inner.manager);
195            let header = oino_session::SessionHeader {
196                version: oino_session::CURRENT_SESSION_VERSION,
197                id: manager.session_id().to_owned(),
198                timestamp: now.clone(),
199                cwd: manager.cwd().display().to_string(),
200                parent_session: None,
201            };
202            let mut lines = vec![
203                FileLine::Header(header)
204                    .to_line()
205                    .map_err(|e| SessionOpError::Failed(e.to_string()))?,
206            ];
207            let mut parent = None;
208            for entry in manager.branch(None) {
209                let mut entry = entry.clone();
210                entry.parent_id = parent.take();
211                parent = Some(entry.id.clone());
212                lines.push(
213                    FileLine::Entry(entry)
214                        .to_line()
215                        .map_err(|e| SessionOpError::Failed(e.to_string()))?,
216                );
217            }
218            lines
219        };
220        let default_name = format!("session-{}.jsonl", now.as_str().replace([':', '.'], "-"));
221        let target = match path {
222            Some(path) if path.is_absolute() => path,
223            Some(path) => self.inner.cwd.join(path),
224            None => self.inner.cwd.join(default_name),
225        };
226        if let Some(dir) = target.parent() {
227            tokio::fs::create_dir_all(dir)
228                .await
229                .map_err(|e| SessionOpError::Failed(e.to_string()))?;
230        }
231        let mut text = lines.join("\n");
232        text.push('\n');
233        tokio::fs::write(&target, text)
234            .await
235            .map_err(|e| SessionOpError::Failed(e.to_string()))?;
236        Ok(target)
237    }
238
239    /// Exports the current branch as a self-contained HTML page (`exportToHtml`
240    /// structure: header, messages, tool calls, collapsed thinking, embedded CSS).
241    ///
242    /// # Errors
243    ///
244    /// Write failures.
245    pub async fn export_html(&self, path: Option<PathBuf>) -> Result<PathBuf, SessionOpError> {
246        let id = self.session_id();
247        let branch: Vec<SessionEntry> = lock(&self.inner.manager)
248            .branch(None)
249            .into_iter()
250            .cloned()
251            .collect();
252        let html = crate::export::render_html(&id, self.session_name().as_deref(), &branch);
253        let short: String = id.chars().take(8).collect();
254        let target =
255            path.unwrap_or_else(|| self.inner.cwd.join(format!("oino-session-{short}.html")));
256        tokio::fs::write(&target, html)
257            .await
258            .map_err(|e| SessionOpError::Failed(e.to_string()))?;
259        Ok(target)
260    }
261}