obol_core/
lib.rs

1//! obol-core: parse agent transcripts and estimate token cost.
2
3pub mod cost;
4pub mod error;
5pub mod model;
6pub mod pricing;
7pub mod transcript;
8
9pub use error::ObolError;
10pub use model::{
11    Approximation, CostEstimate, MessageUsage, ModelCost, PricingSource, Provider, TokenBuckets,
12};
13pub use transcript::Dialect;
14
15use std::path::{Path, PathBuf};
16
17/// Report from a pricing refresh.
18#[derive(Debug, serde::Serialize)]
19pub struct RefreshReport {
20    pub models: usize,
21    pub as_of: String,
22    pub written_to: PathBuf,
23}
24
25/// Resolve the price snapshot. Explicit OBOL_PRICING_DIR wins absolutely; otherwise
26/// pick whichever of {on-disk current.json, embedded} has the newer `as_of`, on-disk
27/// winning ties; embedded is the floor.
28fn resolve_store() -> Result<(pricing::PriceStore, PricingSource), ObolError> {
29    if std::env::var_os("OBOL_PRICING_DIR").is_some() {
30        let store = pricing::PriceStore::load(&pricing::current_path())?;
31        return Ok((store, PricingSource::Local));
32    }
33    let embedded = pricing::embedded()?;
34    let embedded_key = pricing::as_of::sort_key(&embedded.as_of)?;
35    let local_path = pricing::current_path();
36    if local_path.exists() {
37        if let Ok(local) = pricing::PriceStore::load(&local_path) {
38            // Compare parsed stamps, never raw strings; a local snapshot with an
39            // unparseable as_of (pre-validation era) loses to the embedded floor.
40            if pricing::as_of::sort_key(&local.as_of).is_ok_and(|k| k >= embedded_key) {
41                return Ok((local, PricingSource::Local));
42            }
43        }
44    }
45    Ok((embedded, PricingSource::Bundled))
46}
47
48/// Estimate the cost of a transcript file under the given dialect. Loads the active
49/// price snapshot (bundled fallback) and prices the parsed usage.
50pub fn estimate_cost(path: &Path, dialect: Dialect) -> Result<CostEstimate, ObolError> {
51    let (store, source_kind) = resolve_store()?;
52    let bytes = std::fs::read(path)?;
53    let usages = transcript::parse(&bytes, dialect)?;
54    Ok(cost::estimate(&usages, &store, source_kind))
55}
56
57/// Fetch the LiteLLM sheet and write it as the active snapshot. `as_of` is the
58/// caller's stamp — `YYYY-MM-DD` or `YYYY-MM-DDTHH:MM:SSZ` (the library has no
59/// clock) — validated before any network or disk I/O.
60pub fn refresh_pricing_tables(as_of: &str) -> Result<RefreshReport, ObolError> {
61    pricing::as_of::validate(as_of)?;
62    let mut store = pricing::refresh::fetch_litellm(as_of)?; // {litellm: …}
63    let openrouter = pricing::refresh::fetch_openrouter()?;
64    store
65        .namespaces
66        .insert("openrouter".to_string(), openrouter);
67    let models: usize = store.namespaces.values().map(|m| m.len()).sum();
68    let dir = pricing::pricing_dir();
69    store.save(&dir.join(pricing::as_of::archive_file_name(as_of)))?;
70    let current = pricing::current_path();
71    store.save(&current)?;
72    Ok(RefreshReport {
73        models,
74        as_of: as_of.to_string(),
75        written_to: current,
76    })
77}
78
79#[cfg(test)]
80mod api_tests {
81    use super::*;
82
83    #[test]
84    fn estimate_cost_on_bytes_with_missing_tables_errors() {
85        std::env::set_var("OBOL_PRICING_DIR", "/nonexistent/obol-xyz");
86        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
87        std::fs::write(
88            &tmp,
89            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
90        )
91        .unwrap();
92        assert!(matches!(
93            estimate_cost(&tmp, Dialect::Claude),
94            Err(ObolError::PricingTablesMissing(_))
95        ));
96        std::fs::remove_file(&tmp).ok();
97        std::env::remove_var("OBOL_PRICING_DIR");
98    }
99
100    #[test]
101    fn estimate_cost_end_to_end_with_seeded_store() {
102        let dir = std::env::temp_dir().join(format!("obol-api-{}", std::process::id()));
103        std::env::set_var("OBOL_PRICING_DIR", &dir);
104        // seed the store from the sample sheet
105        let store = pricing::refresh::normalize_litellm(
106            include_bytes!("../tests/fixtures/litellm-sample.json"),
107            "2026-06-04",
108        )
109        .unwrap();
110        store.save(&pricing::current_path()).unwrap();
111
112        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
113        std::fs::write(
114            &tmp,
115            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
116        )
117        .unwrap();
118        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
119        assert!(est.total_usd > 0.0);
120        assert_eq!(est.pricing_as_of, "2026-06-04");
121        std::fs::remove_file(&tmp).ok();
122
123        std::fs::remove_dir_all(&dir).ok();
124        std::env::remove_var("OBOL_PRICING_DIR");
125    }
126
127    #[test]
128    fn estimate_cost_from_path_then_detect() {
129        let dir = std::env::temp_dir().join(format!("obol-path-{}", std::process::id()));
130        std::env::set_var("OBOL_PRICING_DIR", &dir);
131        let store = pricing::refresh::normalize_litellm(
132            include_bytes!("../tests/fixtures/litellm-sample.json"),
133            "2026-06-04",
134        )
135        .unwrap();
136        store.save(&pricing::current_path()).unwrap();
137
138        // Write the Claude fixture to a real file, detect dialect, then price.
139        let transcript = dir.join("session.jsonl");
140        std::fs::write(
141            &transcript,
142            include_bytes!("../tests/fixtures/claude-mini.jsonl"),
143        )
144        .unwrap();
145        let bytes = std::fs::read(&transcript).unwrap();
146        let d = transcript::detect(&bytes).unwrap();
147        let est = estimate_cost(&transcript, d).unwrap();
148        assert!(est.total_usd > 0.0);
149
150        std::fs::remove_dir_all(&dir).ok();
151        std::env::remove_var("OBOL_PRICING_DIR");
152    }
153
154    #[test]
155    fn falls_back_to_embedded_when_no_local_snapshot() {
156        // Force "no on-disk snapshot" hermetically: point XDG at an empty dir so
157        // `current_path()` resolves to a nonexistent file and the embedded snapshot
158        // is used. (Setting OBOL_PRICING_DIR instead would take the explicit-override
159        // branch and error PricingTablesMissing rather than fall back to embedded.)
160        let xdg = std::env::temp_dir().join(format!("obol-xdg-{}", std::process::id()));
161        std::fs::create_dir_all(&xdg).unwrap();
162        std::env::remove_var("OBOL_PRICING_DIR");
163        std::env::set_var("XDG_DATA_HOME", &xdg);
164        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
165        std::fs::write(
166            &tmp,
167            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
168        )
169        .unwrap();
170        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
171        assert_eq!(est.pricing_source, crate::model::PricingSource::Bundled);
172        assert!(est.total_usd > 0.0, "embedded snapshot should price claude");
173        std::fs::remove_file(&tmp).ok();
174        std::env::remove_var("XDG_DATA_HOME");
175        std::fs::remove_dir_all(&xdg).ok();
176    }
177
178    #[test]
179    fn explicit_override_uses_local_source() {
180        let dir = std::env::temp_dir().join(format!("obol-resolve-{}", std::process::id()));
181        std::env::set_var("OBOL_PRICING_DIR", &dir);
182        let store = pricing::refresh::normalize_litellm(
183            include_bytes!("../tests/fixtures/litellm-sample.json"),
184            "2099-01-01",
185        )
186        .unwrap();
187        store.save(&pricing::current_path()).unwrap();
188        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
189        std::fs::write(
190            &tmp,
191            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
192        )
193        .unwrap();
194        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
195        assert_eq!(est.pricing_source, crate::model::PricingSource::Local);
196        std::fs::remove_file(&tmp).ok();
197        std::fs::remove_dir_all(&dir).ok();
198        std::env::remove_var("OBOL_PRICING_DIR");
199    }
200
201    #[test]
202    fn kimi_model_surfaces_unpriced_loudly() {
203        std::env::remove_var("OBOL_PRICING_DIR");
204        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
205        std::fs::write(
206            &tmp,
207            include_bytes!("../tests/fixtures/kimi-mini.jsonl").as_slice(),
208        )
209        .unwrap();
210        let est = estimate_cost(&tmp, Dialect::Kimi).unwrap();
211        assert_eq!(est.total_usd, 0.0, "kimi-for-coding is unpriced -> $0");
212        assert!(
213            est.unpriced_models.contains(&"kimi-for-coding".to_string()),
214            "must name the unpriced model: {:?}",
215            est.unpriced_models
216        );
217        std::fs::remove_file(&tmp).ok();
218    }
219
220    #[test]
221    fn local_snapshot_with_invalid_as_of_loses_to_embedded() {
222        // "junk-zzzz" sorts lexicographically above any ISO date, which is exactly
223        // the bug: precedence must be decided by parsed stamps, not raw strings.
224        let xdg = std::env::temp_dir().join(format!("obol-xdg-junk-{}", std::process::id()));
225        std::fs::create_dir_all(&xdg).unwrap();
226        std::env::remove_var("OBOL_PRICING_DIR");
227        std::env::set_var("XDG_DATA_HOME", &xdg);
228        let store = pricing::refresh::normalize_litellm(
229            include_bytes!("../tests/fixtures/litellm-sample.json"),
230            "junk-zzzz",
231        )
232        .unwrap();
233        std::fs::create_dir_all(pricing::pricing_dir()).unwrap();
234        store.save(&pricing::current_path()).unwrap();
235        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
236        std::fs::write(
237            &tmp,
238            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
239        )
240        .unwrap();
241        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
242        assert_eq!(
243            est.pricing_source,
244            crate::model::PricingSource::Bundled,
245            "a junk-stamped local snapshot must not beat the embedded floor"
246        );
247        std::fs::remove_file(&tmp).ok();
248        std::env::remove_var("XDG_DATA_HOME");
249        std::fs::remove_dir_all(&xdg).ok();
250    }
251
252    #[test]
253    fn local_datetime_stamp_beats_embedded_date() {
254        let xdg = std::env::temp_dir().join(format!("obol-xdg-dt-{}", std::process::id()));
255        std::fs::create_dir_all(&xdg).unwrap();
256        std::env::remove_var("OBOL_PRICING_DIR");
257        std::env::set_var("XDG_DATA_HOME", &xdg);
258        let store = pricing::refresh::normalize_litellm(
259            include_bytes!("../tests/fixtures/litellm-sample.json"),
260            "2099-01-01T08:30:00Z",
261        )
262        .unwrap();
263        std::fs::create_dir_all(pricing::pricing_dir()).unwrap();
264        store.save(&pricing::current_path()).unwrap();
265        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
266        std::fs::write(
267            &tmp,
268            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
269        )
270        .unwrap();
271        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
272        assert_eq!(est.pricing_source, crate::model::PricingSource::Local);
273        std::fs::remove_file(&tmp).ok();
274        std::env::remove_var("XDG_DATA_HOME");
275        std::fs::remove_dir_all(&xdg).ok();
276    }
277
278    #[test]
279    fn refresh_rejects_invalid_as_of_before_any_network_or_disk_io() {
280        let xdg = std::env::temp_dir().join(format!("obol-xdg-rej-{}", std::process::id()));
281        std::fs::create_dir_all(&xdg).unwrap();
282        std::env::remove_var("OBOL_PRICING_DIR");
283        std::env::set_var("XDG_DATA_HOME", &xdg);
284        assert!(matches!(
285            refresh_pricing_tables("Apr-2027"),
286            Err(ObolError::InvalidAsOf(_))
287        ));
288        assert!(
289            !pricing::current_path().exists(),
290            "rejected refresh must not write a snapshot"
291        );
292        std::env::remove_var("XDG_DATA_HOME");
293        std::fs::remove_dir_all(&xdg).ok();
294    }
295
296    #[test]
297    fn refresh_report_serializes() {
298        let r = RefreshReport {
299            models: 7,
300            as_of: "2026-06-05".into(),
301            written_to: "/x/current.json".into(),
302        };
303        let v = serde_json::to_value(&r).unwrap();
304        assert_eq!(v["models"], 7);
305        assert_eq!(v["as_of"], "2026-06-05");
306        assert_eq!(v["written_to"], "/x/current.json");
307    }
308}
obol_core/lib.rs

obol_core/
lib.rs
