1use 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#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16pub struct ContextUsage {
17 pub tokens: Option<u64>,
19 pub context_window: u64,
21 pub percent: Option<f64>,
23}
24
25#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(rename_all = "camelCase")]
28pub struct TokenTotals {
29 pub input: u64,
31 pub output: u64,
33 pub cache_read: u64,
35 pub cache_write: u64,
37 pub total: u64,
39}
40
41#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
43#[serde(rename_all = "camelCase")]
44pub struct SessionStats {
45 #[serde(default, skip_serializing_if = "Option::is_none")]
47 pub session_file: Option<String>,
48 pub session_id: String,
50 pub user_messages: u64,
52 pub assistant_messages: u64,
54 pub tool_calls: u64,
56 pub tool_results: u64,
58 pub total_messages: u64,
60 pub tokens: TokenTotals,
62 pub cost: f64,
64 #[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 #[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 #[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 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 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}