Skip to main content

ao_core/
cost_log.rs

1//! Workspace-local usage JSONL (one line per agent turn).
2//!
3//! Companion to `activity_log.rs`: lives at `{workspace}/.ao/usage.jsonl`
4//! and carries aggregated token usage for plugins that don't have a
5//! native JSONL source of their own. The default `Agent::cost_estimate`
6//! reads this file; plugins that override `cost_estimate` (e.g.
7//! `agent-claude-code` reading `~/.claude/projects/**`) ignore it.
8//!
9//! Distinct from `cost_ledger.rs`:
10//! - `cost_ledger` is the **daemon-side** monthly rotation under
11//!   `~/.ao-rs/cost-ledger/YYYY-MM.yaml` — keyed by session id across
12//!   workspaces.
13//! - `cost_log` is the **per-workspace** turn-level log — one file per
14//!   session's worktree.
15//!
16//! Format: newline-delimited JSON, one `UsageLogEntry` per line:
17//!
18//! ```json
19//! {"ts":"2026-04-17T03:07:00Z","input_tokens":500,"output_tokens":200,
20//!  "cache_read_tokens":10,"cache_creation_tokens":5,"cost_usd":0.0034}
21//! ```
22//!
23//! Every field is `#[serde(default)]` so partial writes contribute what
24//! they can. Unknown fields are ignored (forward compat).
25
26use crate::types::CostEstimate;
27use serde::{Deserialize, Serialize};
28use std::io::BufRead;
29use std::path::{Path, PathBuf};
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct UsageLogEntry {
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub ts: Option<String>,
35    #[serde(default)]
36    pub input_tokens: u64,
37    #[serde(default)]
38    pub output_tokens: u64,
39    #[serde(default)]
40    pub cache_read_tokens: u64,
41    #[serde(default)]
42    pub cache_creation_tokens: u64,
43    #[serde(default)]
44    pub cost_usd: f64,
45}
46
47/// Canonical log path: `{workspace}/.ao/usage.jsonl`.
48pub fn usage_log_path(workspace_path: &Path) -> PathBuf {
49    workspace_path.join(".ao").join("usage.jsonl")
50}
51
52/// Aggregate every parseable line in `{workspace}/.ao/usage.jsonl`
53/// into a single `CostEstimate`.
54///
55/// Returns `None` when the file is missing, empty, unreadable, or
56/// aggregates to zero tokens — matching the "no data == no estimate"
57/// rule used by plugins with native parsers.
58pub fn parse_usage_jsonl(workspace_path: &Path) -> Option<CostEstimate> {
59    let path = usage_log_path(workspace_path);
60    let file = std::fs::File::open(&path).ok()?;
61    let reader = std::io::BufReader::new(file);
62
63    let mut input_tokens = 0u64;
64    let mut output_tokens = 0u64;
65    let mut cache_read_tokens = 0u64;
66    let mut cache_creation_tokens = 0u64;
67    let mut cost_usd = 0f64;
68
69    for line in reader.lines().map_while(std::result::Result::ok) {
70        if line.trim().is_empty() {
71            continue;
72        }
73        let Ok(e) = serde_json::from_str::<UsageLogEntry>(&line) else {
74            continue;
75        };
76        input_tokens = input_tokens.saturating_add(e.input_tokens);
77        output_tokens = output_tokens.saturating_add(e.output_tokens);
78        cache_read_tokens = cache_read_tokens.saturating_add(e.cache_read_tokens);
79        cache_creation_tokens = cache_creation_tokens.saturating_add(e.cache_creation_tokens);
80        cost_usd += e.cost_usd;
81    }
82
83    if input_tokens == 0 && output_tokens == 0 {
84        return None;
85    }
86
87    Some(CostEstimate {
88        input_tokens,
89        output_tokens,
90        cache_read_tokens,
91        cache_creation_tokens,
92        cost_usd: if cost_usd > 0.0 { Some(cost_usd) } else { None },
93    })
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use std::io::Write;
100    use std::time::{SystemTime, UNIX_EPOCH};
101
102    fn unique_workspace(label: &str) -> PathBuf {
103        let nanos = SystemTime::now()
104            .duration_since(UNIX_EPOCH)
105            .unwrap()
106            .as_nanos();
107        let p = std::env::temp_dir().join(format!("ao-rs-cost-log-{label}-{nanos}"));
108        std::fs::create_dir_all(&p).unwrap();
109        p
110    }
111
112    fn write_log(workspace: &Path, lines: &[&str]) {
113        let path = usage_log_path(workspace);
114        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
115        let mut f = std::fs::File::create(&path).unwrap();
116        for line in lines {
117            writeln!(f, "{line}").unwrap();
118        }
119    }
120
121    #[test]
122    fn missing_file_returns_none() {
123        let ws = unique_workspace("missing");
124        assert!(parse_usage_jsonl(&ws).is_none());
125    }
126
127    #[test]
128    fn empty_file_returns_none() {
129        let ws = unique_workspace("empty");
130        write_log(&ws, &[]);
131        assert!(parse_usage_jsonl(&ws).is_none());
132    }
133
134    #[test]
135    fn zero_tokens_returns_none() {
136        let ws = unique_workspace("zero");
137        write_log(
138            &ws,
139            &[r#"{"input_tokens":0,"output_tokens":0,"cost_usd":0.0}"#],
140        );
141        assert!(parse_usage_jsonl(&ws).is_none());
142    }
143
144    #[test]
145    fn single_line_round_trip() {
146        let ws = unique_workspace("single");
147        write_log(
148            &ws,
149            &[
150                r#"{"input_tokens":100,"output_tokens":50,"cache_read_tokens":10,"cache_creation_tokens":5,"cost_usd":0.0012}"#,
151            ],
152        );
153        let got = parse_usage_jsonl(&ws).expect("some");
154        assert_eq!(got.input_tokens, 100);
155        assert_eq!(got.output_tokens, 50);
156        assert_eq!(got.cache_read_tokens, 10);
157        assert_eq!(got.cache_creation_tokens, 5);
158        assert!((got.cost_usd.unwrap() - 0.0012).abs() < 1e-9);
159    }
160
161    #[test]
162    fn multi_line_sums_all_fields() {
163        let ws = unique_workspace("multi");
164        write_log(
165            &ws,
166            &[
167                r#"{"input_tokens":100,"output_tokens":50,"cost_usd":0.5}"#,
168                r#"{"input_tokens":200,"output_tokens":75,"cache_read_tokens":4,"cost_usd":0.25}"#,
169                r#"{"input_tokens":50,"output_tokens":25,"cache_creation_tokens":2,"cost_usd":0.125}"#,
170            ],
171        );
172        let got = parse_usage_jsonl(&ws).expect("some");
173        assert_eq!(got.input_tokens, 350);
174        assert_eq!(got.output_tokens, 150);
175        assert_eq!(got.cache_read_tokens, 4);
176        assert_eq!(got.cache_creation_tokens, 2);
177        assert!((got.cost_usd.unwrap() - 0.875).abs() < 1e-9);
178    }
179
180    #[test]
181    fn garbage_lines_are_skipped() {
182        let ws = unique_workspace("garbage");
183        write_log(
184            &ws,
185            &[
186                "not json",
187                r#"{"input_tokens":100,"output_tokens":50}"#,
188                "{",
189                r#"{"input_tokens":10,"output_tokens":5}"#,
190                "",
191            ],
192        );
193        let got = parse_usage_jsonl(&ws).expect("some");
194        assert_eq!(got.input_tokens, 110);
195        assert_eq!(got.output_tokens, 55);
196    }
197
198    #[test]
199    fn unknown_fields_are_ignored() {
200        let ws = unique_workspace("unknown");
201        write_log(
202            &ws,
203            &[r#"{"input_tokens":10,"output_tokens":5,"model":"opus","unknown":"ok"}"#],
204        );
205        let got = parse_usage_jsonl(&ws).expect("some");
206        assert_eq!(got.input_tokens, 10);
207        assert_eq!(got.output_tokens, 5);
208    }
209
210    #[test]
211    fn usage_log_path_shape() {
212        let ws = PathBuf::from("/tmp/ao-ws");
213        assert_eq!(
214            usage_log_path(&ws),
215            PathBuf::from("/tmp/ao-ws/.ao/usage.jsonl")
216        );
217    }
218}