use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::message::{AgentMessage, ContentBlock};
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolStat {
pub name: String,
pub calls: u64,
pub errors: u64,
pub result_bytes: u64,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Cost {
pub micro_usd: u64,
pub unpriced_turns: u64,
}
pub const TOOL_LIST_CAP: usize = 5;
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionStats {
pub turns: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_tokens: u64,
pub cache_write_tokens: u64,
pub cache_turns: u64,
pub tools: Vec<ToolStat>,
pub tool_errors: u64,
pub last_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost: Option<Cost>,
}
pub fn session_stats(messages: &[AgentMessage]) -> SessionStats {
let mut stats = SessionStats::default();
let mut index: HashMap<String, usize> = HashMap::new();
let mut call_name: HashMap<&str, &str> = HashMap::new();
for m in messages {
let AgentMessage::Assistant { content, usage, .. } = m else {
continue;
};
if let Some(u) = usage {
stats.turns += 1;
stats.input_tokens += u.input_tokens;
stats.output_tokens += u.output_tokens;
stats.cache_read_tokens += u.cache_read_tokens.unwrap_or(0);
stats.cache_write_tokens += u.cache_write_tokens.unwrap_or(0);
if u.cache_read_tokens.is_some() || u.cache_write_tokens.is_some() {
stats.cache_turns += 1;
}
stats.last_input_tokens = Some(u.input_tokens);
}
for block in content {
if let ContentBlock::ToolCall { id, name, .. } = block {
call_name.insert(id.as_str(), name.as_str());
let i = match index.get(name) {
Some(&i) => i,
None => {
stats.tools.push(ToolStat {
name: name.clone(),
calls: 0,
errors: 0,
result_bytes: 0,
});
let i = stats.tools.len() - 1;
index.insert(name.clone(), i);
i
}
};
stats.tools[i].calls += 1;
}
}
}
for m in messages {
let AgentMessage::ToolResult {
tool_call_id,
output,
is_error,
..
} = m
else {
continue;
};
let Some(name) = call_name.get(tool_call_id.as_str()) else {
continue;
};
if let Some(&i) = index.get(*name) {
if *is_error {
stats.tools[i].errors += 1;
}
stats.tools[i].result_bytes += output.len() as u64;
}
}
stats.tool_errors = stats.tools.iter().map(|t| t.errors).sum();
stats
.tools
.sort_by(|a, b| b.calls.cmp(&a.calls).then_with(|| a.name.cmp(&b.name)));
stats
}
pub fn aggregate(members: &[SessionStats]) -> SessionStats {
let mut out = SessionStats::default();
let mut index: HashMap<String, usize> = HashMap::new();
for m in members {
out.turns += m.turns;
out.input_tokens += m.input_tokens;
out.output_tokens += m.output_tokens;
out.cache_read_tokens += m.cache_read_tokens;
out.cache_write_tokens += m.cache_write_tokens;
out.cache_turns += m.cache_turns;
out.tool_errors += m.tool_errors;
for t in &m.tools {
let i = match index.get(&t.name) {
Some(&i) => i,
None => {
out.tools.push(ToolStat {
name: t.name.clone(),
calls: 0,
errors: 0,
result_bytes: 0,
});
let i = out.tools.len() - 1;
index.insert(t.name.clone(), i);
i
}
};
out.tools[i].calls += t.calls;
out.tools[i].errors += t.errors;
out.tools[i].result_bytes += t.result_bytes;
}
}
let micro_usd = members
.iter()
.filter_map(|m| m.cost.as_ref())
.map(|c| c.micro_usd)
.sum();
let unpriced_turns = members
.iter()
.map(|m| match &m.cost {
Some(c) => c.unpriced_turns,
None => m.turns, })
.sum();
out.cost = members
.iter()
.any(|m| m.cost.is_some())
.then_some(Cost {
micro_usd,
unpriced_turns,
});
out.tools
.sort_by(|a, b| b.calls.cmp(&a.calls).then_with(|| a.name.cmp(&b.name)));
out
}
impl SessionStats {
pub fn summary(&self) -> String {
if self.turns == 0 {
return "(no usage reported)".to_string();
}
let mut parts = vec![
format!(
"{} turn{}",
self.turns,
if self.turns == 1 { "" } else { "s" }
),
format!("{} in", self.input_tokens),
format!("{} out", self.output_tokens),
];
if self.cache_read_tokens > 0 {
parts.push(format!("{} cache read", self.cache_read_tokens));
}
if self.cache_write_tokens > 0 {
parts.push(format!("{} cache write", self.cache_write_tokens));
}
if let Some(pct) = self.cache_hit_pct() {
parts.push(format!("{pct}% cache"));
}
if let Some(tools) = self.tools_segment() {
parts.push(format!("tools: {tools}"));
}
if self.tool_errors > 0 {
parts.push(format!(
"{} tool error{}",
self.tool_errors,
if self.tool_errors == 1 { "" } else { "s" }
));
}
if let Some(cost) = &self.cost {
parts.push(if cost.unpriced_turns == 0 {
format!("~{}", Self::format_usd(cost.micro_usd))
} else {
format!(
"~{} ({} turn{} unpriced)",
Self::format_usd(cost.micro_usd),
cost.unpriced_turns,
if cost.unpriced_turns == 1 { "" } else { "s" }
)
});
}
format!("usage: {}", parts.join(", "))
}
pub fn cache_hit_pct(&self) -> Option<u8> {
if self.input_tokens == 0 || self.cache_turns == 0 {
return None;
}
let pct =
(self.cache_read_tokens * 100 + self.input_tokens / 2) / self.input_tokens;
Some(pct.min(100) as u8)
}
pub fn tool_error_pct(&self) -> Option<u8> {
let calls: u64 = self.tools.iter().map(|t| t.calls).sum();
if calls == 0 {
return None;
}
let pct = (self.tool_errors * 100 + calls / 2) / calls;
Some(pct.min(100) as u8)
}
fn format_usd(micro_usd: u64) -> String {
format!(
"${}.{:04}",
micro_usd / 1_000_000,
(micro_usd % 1_000_000) / 100
)
}
fn tools_segment(&self) -> Option<String> {
if self.tools.is_empty() {
return None;
}
let shown = self.tools.len().min(TOOL_LIST_CAP);
let mut seg = self.tools[..shown]
.iter()
.map(|t| format!("{}×{}", t.name, t.calls))
.collect::<Vec<_>>()
.join(" ");
if self.tools.len() > TOOL_LIST_CAP {
seg.push_str(&format!(" +{} more", self.tools.len() - TOOL_LIST_CAP));
}
Some(seg)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::{StopReason, Usage};
use serde_json::json;
fn assistant_usage(input: u64, output: u64, read: Option<u64>, write: Option<u64>) -> AgentMessage {
AgentMessage::Assistant {
content: vec![ContentBlock::Text { text: "a".into() }],
stop_reason: StopReason::Stop,
usage: Some(Usage {
input_tokens: input,
output_tokens: output,
cache_read_tokens: read,
cache_write_tokens: write,
}),
model: None,
}
}
fn assistant_tools(names: &[&str]) -> AgentMessage {
AgentMessage::Assistant {
content: names
.iter()
.enumerate()
.map(|(i, n)| ContentBlock::ToolCall {
id: format!("t{i}"),
name: (*n).into(),
arguments: json!({}),
})
.collect(),
stop_reason: StopReason::ToolUse,
usage: None,
model: None,
}
}
fn result(id: &str, name: &str, output: &str, is_error: bool) -> AgentMessage {
AgentMessage::ToolResult {
tool_call_id: id.into(),
name: name.into(),
output: output.into(),
is_error,
}
}
fn stat(name: &str, calls: u64, errors: u64, bytes: u64) -> ToolStat {
ToolStat {
name: name.into(),
calls,
errors,
result_bytes: bytes,
}
}
#[test]
fn empty_is_default() {
assert_eq!(session_stats(&[]), SessionStats::default());
}
#[test]
fn counts_usage_only_on_reported_turns() {
let messages = vec![
AgentMessage::user_text("q"),
assistant_usage(10, 20, Some(3), None),
AgentMessage::Assistant {
content: vec![ContentBlock::Text { text: "x".into() }],
stop_reason: StopReason::Stop,
usage: None,
model: None,
},
assistant_usage(30, 40, None, Some(5)),
];
assert_eq!(
session_stats(&messages),
SessionStats {
turns: 2,
input_tokens: 40,
output_tokens: 60,
cache_read_tokens: 3,
cache_write_tokens: 5,
cache_turns: 2,
tools: vec![],
tool_errors: 0,
last_input_tokens: Some(30),
cost: None,
}
);
}
#[test]
fn tallies_tool_calls_by_name() {
let stats = session_stats(&[assistant_tools(&["read", "read", "grep"])]);
assert_eq!(stats.tools, vec![stat("read", 2, 0, 0), stat("grep", 1, 0, 0)]);
assert_eq!(stats.turns, 0);
}
#[test]
fn tools_sorted_count_desc_then_name_asc() {
let stats = session_stats(&[assistant_tools(&["zzz", "aaa", "aaa", "bbb"])]);
assert_eq!(
stats.tools,
vec![stat("aaa", 2, 0, 0), stat("bbb", 1, 0, 0), stat("zzz", 1, 0, 0)]
);
}
#[test]
fn last_input_tokens_is_most_recent() {
let messages = vec![
assistant_usage(10, 20, None, None),
AgentMessage::user_text("q"),
assistant_usage(40, 3, None, None),
assistant_tools(&["read"]), ];
assert_eq!(session_stats(&messages).last_input_tokens, Some(40));
assert_eq!(session_stats(&[]).last_input_tokens, None);
assert_eq!(
session_stats(&[AgentMessage::user_text("q")]).last_input_tokens,
None
);
}
#[test]
fn summary_empty_and_full_line() {
assert_eq!(SessionStats::default().summary(), "(no usage reported)");
let messages = vec![
assistant_usage(40, 60, Some(3), Some(5)),
assistant_tools(&["read"]),
AgentMessage::user_text("q"),
assistant_usage(0, 0, None, None),
];
assert_eq!(
session_stats(&messages).summary(),
"usage: 2 turns, 40 in, 60 out, 3 cache read, 5 cache write, 8% cache, tools: read×1"
);
let one = SessionStats {
turns: 1,
input_tokens: 10,
output_tokens: 20,
..Default::default()
};
assert_eq!(one.summary(), "usage: 1 turn, 10 in, 20 out");
}
#[test]
fn summary_caps_tools_and_says_more() {
let stats = SessionStats {
turns: 1,
input_tokens: 1,
output_tokens: 1,
tools: vec![
stat("a", 1, 0, 0),
stat("b", 1, 0, 0),
stat("c", 1, 0, 0),
stat("d", 1, 0, 0),
stat("e", 1, 0, 0),
stat("f", 1, 0, 0),
stat("g", 1, 0, 0),
],
last_input_tokens: Some(1),
..Default::default()
};
assert_eq!(
stats.summary(),
"usage: 1 turn, 1 in, 1 out, tools: a×1 b×1 c×1 d×1 e×1 +2 more"
);
}
#[test]
fn summary_omits_empty_tools_segment() {
let stats = SessionStats {
turns: 1,
input_tokens: 10,
output_tokens: 20,
..Default::default()
};
let line = stats.summary();
assert_eq!(line, "usage: 1 turn, 10 in, 20 out");
assert!(!line.contains("tools:"));
}
#[test]
fn tool_errors_pair_by_call_id() {
let messages = vec![
assistant_tools(&["read"]), result("t0", "grep", "boom", true), ];
let stats = session_stats(&messages);
assert_eq!(stats.tools, vec![stat("read", 1, 1, 4)]);
assert_eq!(stats.tool_errors, 1);
}
#[test]
fn a_call_without_a_result_counts_as_a_call_not_an_error() {
let stats = session_stats(&[assistant_tools(&["read"])]);
assert_eq!(stats.tools, vec![stat("read", 1, 0, 0)]);
assert_eq!(stats.tool_errors, 0);
}
#[test]
fn a_synthesized_abort_result_is_counted_as_an_error() {
let messages = vec![
assistant_tools(&["read"]),
result("t0", "read", "aborted before execution", true),
];
let stats = session_stats(&messages);
assert_eq!(stats.tools[0].errors, 1);
assert_eq!(stats.tool_errors, 1);
}
#[test]
fn result_bytes_sum_per_tool() {
let messages = vec![
assistant_tools(&["read", "read"]), result("t0", "read", "abc", false), result("t1", "read", "de", false), result("gone", "read", "ignored", false), ];
let stats = session_stats(&messages);
assert_eq!(stats.tools, vec![stat("read", 2, 0, 5)]);
}
#[test]
fn cache_hit_pct_is_none_without_cache_reporting() {
let stats = session_stats(&[assistant_usage(100, 10, None, None)]);
assert_eq!(stats.cache_turns, 0);
assert_eq!(stats.cache_hit_pct(), None);
assert!(!stats.summary().contains("% cache"));
}
#[test]
fn cache_hit_pct_rounds_to_nearest_percent() {
assert_eq!(
session_stats(&[assistant_usage(3, 0, Some(1), None)]).cache_hit_pct(),
Some(33)
);
assert_eq!(
session_stats(&[assistant_usage(3, 0, Some(2), None)]).cache_hit_pct(),
Some(67)
);
assert_eq!(
session_stats(&[assistant_usage(3, 0, Some(3), None)]).cache_hit_pct(),
Some(100)
);
assert_eq!(
session_stats(&[assistant_usage(3, 0, Some(6), None)]).cache_hit_pct(),
Some(100)
);
}
#[test]
fn tool_error_pct_is_none_with_no_calls() {
let stats = session_stats(&[assistant_usage(1, 1, None, None)]);
assert_eq!(stats.tool_error_pct(), None);
assert!(!stats.summary().contains("tool error"));
}
#[test]
fn aggregate_sums_two_members_and_sorts_tools() {
let a = SessionStats {
turns: 1,
input_tokens: 10,
output_tokens: 20,
cache_turns: 1,
tools: vec![stat("read", 2, 0, 5)],
cost: Some(Cost {
micro_usd: 4_000,
unpriced_turns: 0,
}),
last_input_tokens: Some(10),
..Default::default()
};
let b = SessionStats {
turns: 2,
input_tokens: 30,
output_tokens: 40,
tools: vec![stat("read", 1, 1, 7), stat("grep", 3, 0, 2)],
tool_errors: 1,
..Default::default()
};
let total = aggregate(&[a, b]);
assert_eq!(total.turns, 3);
assert_eq!(total.input_tokens, 40);
assert_eq!(total.output_tokens, 60);
assert_eq!(total.cache_turns, 1);
assert_eq!(total.tool_errors, 1);
assert_eq!(
total.tools,
vec![stat("grep", 3, 0, 2), stat("read", 3, 1, 12)]
);
assert_eq!(
total.cost,
Some(Cost {
micro_usd: 4_000,
unpriced_turns: 2,
})
);
assert_eq!(total.last_input_tokens, None);
}
#[test]
fn summary_reports_cost_and_unpriced_turns() {
let priced = SessionStats {
turns: 1,
input_tokens: 1,
output_tokens: 1,
cost: Some(Cost {
micro_usd: 4_000,
unpriced_turns: 0,
}),
..Default::default()
};
assert!(
priced.summary().ends_with("~$0.0040"),
"{}",
priced.summary()
);
let partial = SessionStats {
turns: 1,
input_tokens: 1,
output_tokens: 1,
cost: Some(Cost {
micro_usd: 4_000,
unpriced_turns: 1,
}),
..Default::default()
};
assert!(
partial.summary().ends_with("~$0.0040 (1 turn unpriced)"),
"{}",
partial.summary()
);
}
}