use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl TokenUsage {
pub fn total(&self) -> u32 {
self.input_tokens.saturating_add(self.output_tokens)
}
pub fn saturating_add(self, other: Self) -> Self {
Self {
input_tokens: self.input_tokens.saturating_add(other.input_tokens),
output_tokens: self.output_tokens.saturating_add(other.output_tokens),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TurnEndReason {
Completed,
NarrationCapExhausted,
NarrationFinalRound,
RoundCap,
Empty,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TurnMetrics {
pub elapsed_ms: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<TokenUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_usd: Option<f64>,
pub model_id: String,
pub endpoint: String,
#[serde(default, skip_serializing_if = "is_zero_u32")]
pub hallucinations: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub end_reason: Option<TurnEndReason>,
}
fn is_zero_u32(n: &u32) -> bool {
*n == 0
}
impl TurnMetrics {
pub fn display_line(&self) -> String {
let elapsed = if self.elapsed_ms >= 1000 {
format!("{:.1}s", self.elapsed_ms as f64 / 1000.0)
} else {
format!("{}ms", self.elapsed_ms)
};
let token_part = match self.usage {
Some(u) => format!(
"{} in / {} out",
fmt_count(u.input_tokens),
fmt_count(u.output_tokens)
),
None => "(tokens unavailable)".into(),
};
let cost_part = match self.cost_usd {
Some(c) if c < f64::EPSILON => "free (local)".into(),
Some(c) if c < 0.001 => format!("~${c:.5}"),
Some(c) if c < 0.01 => format!("~${c:.4}"),
Some(c) => format!("~${c:.4}"),
None if self.usage.is_some() => "free (local)".into(),
None => String::new(),
};
let base = if cost_part.is_empty() {
format!("{elapsed} · {token_part}")
} else {
format!("{elapsed} · {token_part} · {cost_part}")
};
let base = if self.hallucinations > 0 {
format!(
"{base} · ⚠ {} hallucination(s) corrected",
self.hallucinations
)
} else {
base
};
match self.end_reason {
Some(TurnEndReason::NarrationCapExhausted) => {
format!("{base} · ⚠ ended on narration (rescue budget spent)")
}
Some(TurnEndReason::NarrationFinalRound) => {
format!("{base} · ⚠ ended on narration (final round)")
}
Some(TurnEndReason::RoundCap) => format!("{base} · round cap"),
Some(TurnEndReason::Empty) => format!("{base} · ⚠ empty response"),
Some(TurnEndReason::Completed) | None => base,
}
}
pub fn prometheus_labels(&self) -> [(&'static str, String); 2] {
[
("model", self.model_id.clone()),
("endpoint", self.endpoint.clone()),
]
}
pub fn append_to_log(&self, path: &std::path::Path) {
let _ = append_jsonl(self, path);
}
pub fn append_to_log_with_policy(&self, path: &std::path::Path, policy: &crate::LogConfig) {
let _ = append_jsonl(self, path);
let _ = rotate_log(path, policy);
}
}
fn fmt_count(n: u32) -> String {
let s = n.to_string();
let mut out = String::with_capacity(s.len() + s.len() / 3);
for (i, c) in s.chars().rev().enumerate() {
if i > 0 && i % 3 == 0 {
out.push(',');
}
out.push(c);
}
out.chars().rev().collect()
}
fn append_jsonl(metrics: &TurnMetrics, path: &std::path::Path) -> std::io::Result<()> {
use std::io::Write as _;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut line = serde_json::to_string(metrics).map_err(std::io::Error::other)?;
line.push('\n');
let mut f = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)?;
f.write_all(line.as_bytes())
}
fn rotate_log(path: &std::path::Path, policy: &crate::LogConfig) -> std::io::Result<()> {
if policy.max_sessions > 0 {
let content = std::fs::read_to_string(path)?;
let lines: Vec<&str> = content.lines().collect();
if lines.len() > policy.max_sessions {
let kept = lines[lines.len() - policy.max_sessions..].join("\n");
std::fs::write(path, format!("{kept}\n"))?;
}
}
if policy.max_size_mb > 0 {
let meta = std::fs::metadata(path)?;
let limit_bytes = policy.max_size_mb * 1024 * 1024;
if meta.len() > limit_bytes {
for i in (1..=policy.keep_rotated).rev() {
let older = path.with_extension(format!("jsonl.{i}"));
let newer = if i == 1 {
path.to_path_buf()
} else {
path.with_extension(format!("jsonl.{}", i - 1))
};
if newer.exists() {
if i == policy.keep_rotated && older.exists() {
let _ = std::fs::remove_file(&older);
}
let _ = std::fs::rename(&newer, &older);
}
}
std::fs::File::create(path)?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn metrics(elapsed_ms: u64, in_tok: u32, out_tok: u32, cost: Option<f64>) -> TurnMetrics {
TurnMetrics {
elapsed_ms,
usage: Some(TokenUsage {
input_tokens: in_tok,
output_tokens: out_tok,
}),
cost_usd: cost,
model_id: "gemma4:e2b".into(),
endpoint: "http://REDACTED-HOST:11434".into(),
..Default::default()
}
}
#[test]
fn display_free_local() {
let m = metrics(3200, 847, 312, Some(0.0));
let line = m.display_line();
assert!(line.starts_with("3.2s"), "got: {line}");
assert!(line.contains("847"), "got: {line}");
assert!(line.contains("312"), "got: {line}");
assert!(line.contains("free (local)"), "got: {line}");
}
#[test]
fn display_with_cost() {
let m = metrics(8700, 1204, 892, Some(0.0041));
let line = m.display_line();
assert!(line.contains("$0.0041"), "got: {line}");
assert!(line.contains("1,204"), "got: {line}");
}
#[test]
fn display_tokens_unavailable() {
let m = TurnMetrics {
elapsed_ms: 5100,
usage: None,
cost_usd: None,
model_id: "gpt-4o".into(),
endpoint: "https://api.openai.com".into(),
..Default::default()
};
let line = m.display_line();
assert!(line.contains("tokens unavailable"), "got: {line}");
}
#[test]
fn display_end_reason_flags_anomalous_endings() {
let mut m = metrics(3200, 100, 50, Some(0.0));
assert!(!m.display_line().contains("narration"), "unset stays quiet");
m.end_reason = Some(TurnEndReason::Completed);
assert!(
!m.display_line().contains("narration"),
"a normal completion stays quiet"
);
m.end_reason = Some(TurnEndReason::NarrationCapExhausted);
assert!(m
.display_line()
.contains("ended on narration (rescue budget spent)"));
m.end_reason = Some(TurnEndReason::NarrationFinalRound);
assert!(m
.display_line()
.contains("ended on narration (final round)"));
m.end_reason = Some(TurnEndReason::RoundCap);
assert!(m.display_line().contains("round cap"));
m.end_reason = Some(TurnEndReason::Empty);
assert!(m.display_line().contains("empty response"));
}
#[test]
fn end_reason_serializes_snake_case_and_skips_when_absent() {
let mut m = metrics(1, 1, 1, None);
let j = serde_json::to_string(&m).unwrap();
assert!(!j.contains("end_reason"), "{j}");
m.end_reason = Some(TurnEndReason::NarrationCapExhausted);
let j = serde_json::to_string(&m).unwrap();
assert!(
j.contains("\"end_reason\":\"narration_cap_exhausted\""),
"{j}"
);
}
#[test]
fn display_milliseconds_under_one_second() {
let m = metrics(850, 100, 50, None);
assert!(
m.display_line().starts_with("850ms"),
"got: {}",
m.display_line()
);
}
#[test]
fn fmt_count_thousands() {
assert_eq!(fmt_count(1000), "1,000");
assert_eq!(fmt_count(1234567), "1,234,567");
assert_eq!(fmt_count(42), "42");
}
#[test]
fn token_usage_total() {
let u = TokenUsage {
input_tokens: 300,
output_tokens: 150,
};
assert_eq!(u.total(), 450);
}
#[test]
fn token_usage_saturating_add() {
let a = TokenUsage {
input_tokens: 100,
output_tokens: 50,
};
let b = TokenUsage {
input_tokens: 200,
output_tokens: 75,
};
let sum = a.saturating_add(b);
assert_eq!(sum.input_tokens, 300);
assert_eq!(sum.output_tokens, 125);
let big = TokenUsage {
input_tokens: u32::MAX,
output_tokens: u32::MAX,
};
let sat = big.saturating_add(b);
assert_eq!(sat.input_tokens, u32::MAX);
}
#[test]
fn display_with_hallucinations() {
let mut m = metrics(3200, 847, 312, Some(0.0));
m.hallucinations = 2;
let line = m.display_line();
assert!(line.contains("2 hallucination(s) corrected"), "got: {line}");
}
#[test]
fn display_no_hallucinations_omits_warning() {
let m = metrics(3200, 847, 312, Some(0.0));
assert_eq!(m.hallucinations, 0);
assert!(
!m.display_line().contains("hallucination"),
"zero hallucinations must not appear in display"
);
}
#[test]
fn hallucinations_zero_skipped_in_json() {
let m = metrics(1000, 10, 5, Some(0.0));
let json = serde_json::to_string(&m).unwrap();
assert!(
!json.contains("hallucination"),
"zero hallucinations must be omitted from JSON"
);
}
#[test]
fn hallucinations_nonzero_in_json() {
let mut m = metrics(1000, 10, 5, Some(0.0));
m.hallucinations = 3;
let json = serde_json::to_string(&m).unwrap();
assert!(json.contains("\"hallucinations\":3"), "got: {json}");
}
#[test]
fn append_to_log_creates_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("usage.jsonl");
let m = metrics(1000, 10, 5, Some(0.0));
m.append_to_log(&path);
let content = std::fs::read_to_string(&path).unwrap();
assert!(content.contains("gemma4:e2b"));
assert!(content.ends_with('\n'));
m.append_to_log(&path);
let raw = std::fs::read_to_string(&path).unwrap();
let lines: Vec<_> = raw.lines().collect();
assert_eq!(lines.len(), 2);
}
#[test]
fn rotation_session_limit_trims_oldest() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("usage.jsonl");
let m = metrics(1000, 10, 5, Some(0.0));
for _ in 0..10 {
m.append_to_log(&path);
}
let policy = crate::LogConfig {
max_sessions: 7,
..Default::default()
};
rotate_log(&path, &policy).unwrap();
let n = std::fs::read_to_string(&path).unwrap().lines().count();
assert_eq!(n, 7, "should keep exactly max_sessions entries");
}
#[test]
fn rotation_session_limit_noop_when_under_cap() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("usage.jsonl");
let m = metrics(1000, 10, 5, Some(0.0));
for _ in 0..5 {
m.append_to_log(&path);
}
let policy = crate::LogConfig {
max_sessions: 7,
..Default::default()
};
rotate_log(&path, &policy).unwrap();
let lines = std::fs::read_to_string(&path).unwrap().lines().count();
assert_eq!(lines, 5, "under cap — no entries should be dropped");
}
#[test]
fn rotation_size_limit_renames_and_resets() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("usage.jsonl");
let m = metrics(1000, 10, 5, Some(0.0));
m.append_to_log(&path);
let policy = crate::LogConfig {
max_sessions: 0,
max_size_mb: 0, keep_rotated: 2,
..Default::default()
};
let policy_tiny = crate::LogConfig {
max_size_mb: 0, ..policy
};
rotate_log(&path, &policy_tiny).unwrap();
assert!(path.exists(), "file must still exist when size limit is 0");
}
#[test]
fn log_config_default_is_7_sessions() {
let cfg = crate::LogConfig::default();
assert_eq!(cfg.max_sessions, 7);
assert_eq!(cfg.max_size_mb, 0);
assert_eq!(cfg.max_age_days, 0);
assert_eq!(cfg.keep_rotated, 3);
}
}