This commit is contained in:
+323
@@ -0,0 +1,323 @@
|
||||
//! Cost tracking (issue #49): collect token usage exposed by agents, price
|
||||
//! it against the configured model sheet, and aggregate per agent/session.
|
||||
//! No runtime dependency is added — everything is std + the existing config.
|
||||
|
||||
use crate::config::{self, CostModel};
|
||||
use crate::events::{self, Event, EventKind};
|
||||
use anyhow::Result;
|
||||
use serde::Serialize;
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Token usage collected from an agent.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize)]
|
||||
pub struct Usage {
|
||||
pub tokens_in: u64,
|
||||
pub tokens_out: u64,
|
||||
}
|
||||
|
||||
/// One aggregated cost row (per agent, or global totals).
|
||||
#[derive(Debug, Clone, Default, Serialize)]
|
||||
pub struct CostStats {
|
||||
pub agent: String,
|
||||
pub sessions: u64,
|
||||
pub tokens_in: u64,
|
||||
pub tokens_out: u64,
|
||||
pub cost_usd: f64,
|
||||
}
|
||||
|
||||
/// Price the usage against a model of the sheet (issue #49).
|
||||
/// Unknown models fall back to the first model of the sheet.
|
||||
pub fn estimate_cost(
|
||||
model: Option<&str>,
|
||||
usage: Usage,
|
||||
sheet: &BTreeMap<String, CostModel>,
|
||||
) -> f64 {
|
||||
let m = model
|
||||
.and_then(|k| sheet.get(k))
|
||||
.or_else(|| sheet.values().next());
|
||||
match m {
|
||||
Some(cm) => {
|
||||
usage.tokens_in as f64 * cm.input_per_mtok / 1_000_000.0
|
||||
+ usage.tokens_out as f64 * cm.output_per_mtok / 1_000_000.0
|
||||
}
|
||||
None => 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Collect the usage an agent exposed for its latest session (issue #49).
|
||||
///
|
||||
/// Supported sources today:
|
||||
/// - Claude Code: `~/.claude/projects/*/<session>.jsonl` — the last JSON
|
||||
/// line carrying an `usage` object with `input_tokens`/`output_tokens`.
|
||||
///
|
||||
/// Returns None when the agent exposes nothing (safe, silent).
|
||||
pub fn collect_usage(agent: &str, cwd: Option<&str>) -> Option<Usage> {
|
||||
match agent {
|
||||
"claude-code" => claude_usage(cwd),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Claude Code stores one JSONL per session under ~/.claude/projects/<enc>/.
|
||||
/// The final line with a "usage" object holds the cumulative token counts.
|
||||
fn claude_usage(cwd: Option<&str>) -> Option<Usage> {
|
||||
let root = config::home_dir()?.join(".claude/projects");
|
||||
let mut candidates: Vec<(String, PathBuf)> = Vec::new();
|
||||
for entry in std::fs::read_dir(&root).ok()?.flatten() {
|
||||
let dir = entry.path();
|
||||
if !dir.is_dir() {
|
||||
continue;
|
||||
}
|
||||
for f in std::fs::read_dir(&dir).ok()?.flatten() {
|
||||
let p = f.path();
|
||||
if p.extension().and_then(|e| e.to_str()) == Some("jsonl") {
|
||||
if let Ok(mtime) = f.metadata().and_then(|m| m.modified()) {
|
||||
candidates.push((mtime_key(mtime), p));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
candidates.sort();
|
||||
let last = candidates.last()?.1.clone();
|
||||
parse_claude_file(&last, cwd)
|
||||
}
|
||||
|
||||
fn mtime_key(t: std::time::SystemTime) -> String {
|
||||
format!(
|
||||
"{:?}",
|
||||
t.duration_since(std::time::UNIX_EPOCH).unwrap_or_default()
|
||||
)
|
||||
}
|
||||
|
||||
/// Parse a Claude Code session file: the last line whose "usage" object has
|
||||
/// input/output token counts. Filters by cwd when given (the file embeds it
|
||||
/// as `cwd` on summary lines).
|
||||
fn parse_claude_file(path: &std::path::Path, cwd: Option<&str>) -> Option<Usage> {
|
||||
let text = std::fs::read_to_string(path).ok()?;
|
||||
let mut usage = None;
|
||||
for line in text.lines() {
|
||||
let v: serde_json::Value = serde_json::from_str(line).ok()?;
|
||||
if v.get("usage").is_none() {
|
||||
continue;
|
||||
}
|
||||
if let Some(want) = cwd {
|
||||
if v.get("cwd").and_then(|c| c.as_str()) != Some(want) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let u = v.get("usage")?;
|
||||
let tin = u.get("input_tokens").and_then(|n| n.as_u64())?;
|
||||
let tout = u.get("output_tokens").and_then(|n| n.as_u64())?;
|
||||
usage = Some(Usage {
|
||||
tokens_in: tin,
|
||||
tokens_out: tout,
|
||||
});
|
||||
}
|
||||
usage
|
||||
}
|
||||
|
||||
/// Record a session's usage: persist it on the session record and emit a
|
||||
/// `cost` event (axe 8 journal). Called when an agent session finishes.
|
||||
pub fn record_session_usage(
|
||||
app: &crate::app::App,
|
||||
agent: &str,
|
||||
session_id: &str,
|
||||
cwd: Option<&str>,
|
||||
) -> Result<()> {
|
||||
let Some(usage) = collect_usage(agent, cwd) else {
|
||||
return Ok(());
|
||||
};
|
||||
let model = app
|
||||
.catalog
|
||||
.resolve(agent)
|
||||
.and_then(|a| a.cost_model.as_deref());
|
||||
let sheet = app
|
||||
.config
|
||||
.settings
|
||||
.cost_models
|
||||
.clone()
|
||||
.unwrap_or_else(config::default_cost_models);
|
||||
let cost = estimate_cost(model, usage, &sheet);
|
||||
|
||||
// Persist on the session record.
|
||||
let mut sf = crate::sessions::load(app)?;
|
||||
if let Some(r) = sf.sessions.get_mut(session_id) {
|
||||
r.tokens_in = Some(usage.tokens_in);
|
||||
r.tokens_out = Some(usage.tokens_out);
|
||||
r.cost_usd = Some(cost);
|
||||
crate::sessions::save(app, &sf)?;
|
||||
}
|
||||
|
||||
// Journal a dedicated event (stable --json contract, additive fields).
|
||||
app.emit(
|
||||
&Event::now(EventKind::Cost)
|
||||
.with_agent(agent.to_string())
|
||||
.with_session(session_id)
|
||||
.with_usage(usage.tokens_in, usage.tokens_out)
|
||||
.with_cost(cost),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Aggregate cost events into per-agent rows (issue #49). `--json` output
|
||||
/// stays stable: agent, sessions, tokens_in, tokens_out, cost_usd.
|
||||
pub fn aggregate(
|
||||
app: &crate::app::App,
|
||||
period: Option<&str>,
|
||||
) -> Result<Vec<CostStats>> {
|
||||
let cut = crate::commands::stats_cmd::cutoff(period)?;
|
||||
let events = events::read_events(&app.events_dir(), 0);
|
||||
let mut map: BTreeMap<String, CostStats> = BTreeMap::new();
|
||||
for ev in events {
|
||||
if ev.kind != EventKind::Cost {
|
||||
continue;
|
||||
}
|
||||
if let Some(cut) = &cut {
|
||||
if let Ok(ts) = chrono::DateTime::parse_from_rfc3339(&ev.ts) {
|
||||
if ts.with_timezone(&chrono::Utc) < *cut {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
let agent = ev.agent.clone().unwrap_or_else(|| "unknown".to_string());
|
||||
let row = map.entry(agent.clone()).or_default();
|
||||
row.agent = agent;
|
||||
row.sessions += 1;
|
||||
row.tokens_in += ev.tokens_in.unwrap_or(0);
|
||||
row.tokens_out += ev.tokens_out.unwrap_or(0);
|
||||
row.cost_usd += ev.cost_usd.unwrap_or(0.0);
|
||||
}
|
||||
Ok(map.into_values().collect())
|
||||
}
|
||||
|
||||
/// Total across every row (used by the weekly/monthly report).
|
||||
pub fn totals(rows: &[CostStats]) -> CostStats {
|
||||
let mut t = CostStats {
|
||||
agent: "total".to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
for r in rows {
|
||||
t.sessions += r.sessions;
|
||||
t.tokens_in += r.tokens_in;
|
||||
t.tokens_out += r.tokens_out;
|
||||
t.cost_usd += r.cost_usd;
|
||||
}
|
||||
t
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn estimate_cost_uses_model_prices() {
|
||||
let sheet = config::default_cost_models();
|
||||
// 1M input @ $3 + 500k output @ $15 = 3 + 7.5 = $10.5
|
||||
let u = Usage {
|
||||
tokens_in: 1_000_000,
|
||||
tokens_out: 500_000,
|
||||
};
|
||||
let c = estimate_cost(Some("claude-sonnet-4"), u, &sheet);
|
||||
assert!((c - 10.5).abs() < 1e-9, "got {c}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_cost_falls_back_to_first_model() {
|
||||
let sheet = config::default_cost_models();
|
||||
let c = estimate_cost(Some("nope-model"), Usage { tokens_in: 0, tokens_out: 0 }, &sheet);
|
||||
assert_eq!(c, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_cost_unknown_sheet_is_zero() {
|
||||
let sheet = BTreeMap::new();
|
||||
let c = estimate_cost(None, Usage { tokens_in: 1000, tokens_out: 1000 }, &sheet);
|
||||
assert_eq!(c, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_claude_file_extracts_last_usage() {
|
||||
let dir = std::env::temp_dir().join(format!("am-costs-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let f = dir.join("session.jsonl");
|
||||
std::fs::write(
|
||||
&f,
|
||||
concat!(
|
||||
"{\"type\":\"assistant\",\"message\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":5}},\"cwd\":\"/x\"}\n",
|
||||
"{\"type\":\"summary\",\"usage\":{\"input_tokens\":100,\"output_tokens\":40},\"cwd\":\"/x\"}\n",
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
let u = parse_claude_file(&f, Some("/x")).unwrap();
|
||||
assert_eq!(u.tokens_in, 100);
|
||||
assert_eq!(u.tokens_out, 40);
|
||||
// Wrong cwd → filtered out.
|
||||
assert!(parse_claude_file(&f, Some("/other")).is_none());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_sums_per_agent() {
|
||||
let app = test_app("costs-agg");
|
||||
app.emit(
|
||||
&Event::now(EventKind::Cost)
|
||||
.with_agent("a".to_string())
|
||||
.with_usage(10, 5)
|
||||
.with_cost(0.5),
|
||||
);
|
||||
app.emit(
|
||||
&Event::now(EventKind::Cost)
|
||||
.with_agent("a".to_string())
|
||||
.with_usage(20, 10)
|
||||
.with_cost(1.0),
|
||||
);
|
||||
app.emit(
|
||||
&Event::now(EventKind::Cost)
|
||||
.with_agent("b".to_string())
|
||||
.with_usage(100, 50)
|
||||
.with_cost(2.0),
|
||||
);
|
||||
let rows = aggregate(&app, None).unwrap();
|
||||
assert_eq!(rows.len(), 2);
|
||||
let a = rows.iter().find(|r| r.agent == "a").unwrap();
|
||||
assert_eq!(a.sessions, 2);
|
||||
assert_eq!(a.tokens_in, 30);
|
||||
assert_eq!(a.tokens_out, 15);
|
||||
assert!((a.cost_usd - 1.5).abs() < 1e-9);
|
||||
let t = totals(&rows);
|
||||
assert!((t.cost_usd - 3.5).abs() < 1e-9);
|
||||
assert_eq!(t.sessions, 3);
|
||||
}
|
||||
|
||||
fn test_app(tag: &str) -> crate::app::App {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let cfg = dir.path().join("config.yaml");
|
||||
std::fs::write(
|
||||
&cfg,
|
||||
concat!(
|
||||
"version: \"1.0\"\n",
|
||||
"settings:\n",
|
||||
" auto_install_deps: false\n",
|
||||
" confirm_before_run: false\n",
|
||||
"agents: []\n",
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
use clap::Parser;
|
||||
let cli = crate::cli::Cli::parse_from([
|
||||
"am",
|
||||
"--config",
|
||||
cfg.to_str().unwrap(),
|
||||
"list",
|
||||
]);
|
||||
let mut app = crate::app::App::from_cli(cli).expect("app should build");
|
||||
app.paths.state_file = dir.path().join(format!("state-{tag}.json"));
|
||||
app.paths.log_dir = dir.path().join("logs");
|
||||
app
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user