324 lines
10 KiB
Rust
324 lines
10 KiB
Rust
//! 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
|
|
}
|
|
}
|