1use 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
47pub fn usage_log_path(workspace_path: &Path) -> PathBuf {
49 workspace_path.join(".ao").join("usage.jsonl")
50}
51
52pub 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}