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 local_path = pricing::current_path();
35    if local_path.exists() {
36        if let Ok(local) = pricing::PriceStore::load(&local_path) {
37            if local.as_of >= embedded.as_of {
38                return Ok((local, PricingSource::Local));
39            }
40        }
41    }
42    Ok((embedded, PricingSource::Bundled))
43}
44
45/// Estimate the cost of a transcript file under the given dialect. Loads the active
46/// price snapshot (bundled fallback) and prices the parsed usage.
47pub fn estimate_cost(path: &Path, dialect: Dialect) -> Result<CostEstimate, ObolError> {
48    let (store, source_kind) = resolve_store()?;
49    let bytes = std::fs::read(path)?;
50    let usages = transcript::parse(&bytes, dialect)?;
51    Ok(cost::estimate(&usages, &store, source_kind))
52}
53
54/// Fetch the LiteLLM sheet and write it as the active snapshot. `as_of` is the
55/// caller's date string (the library has no clock).
56pub fn refresh_pricing_tables(as_of: &str) -> Result<RefreshReport, ObolError> {
57    let mut store = pricing::refresh::fetch_litellm(as_of)?; // {litellm: …}
58    let openrouter = pricing::refresh::fetch_openrouter()?;
59    store
60        .namespaces
61        .insert("openrouter".to_string(), openrouter);
62    let models: usize = store.namespaces.values().map(|m| m.len()).sum();
63    let dir = pricing::pricing_dir();
64    store.save(&dir.join(format!("prices-{as_of}.json")))?;
65    let current = pricing::current_path();
66    store.save(&current)?;
67    Ok(RefreshReport {
68        models,
69        as_of: as_of.to_string(),
70        written_to: current,
71    })
72}
73
74#[cfg(test)]
75mod api_tests {
76    use super::*;
77
78    #[test]
79    fn estimate_cost_on_bytes_with_missing_tables_errors() {
80        std::env::set_var("OBOL_PRICING_DIR", "/nonexistent/obol-xyz");
81        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
82        std::fs::write(
83            &tmp,
84            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
85        )
86        .unwrap();
87        assert!(matches!(
88            estimate_cost(&tmp, Dialect::Claude),
89            Err(ObolError::PricingTablesMissing(_))
90        ));
91        std::fs::remove_file(&tmp).ok();
92        std::env::remove_var("OBOL_PRICING_DIR");
93    }
94
95    #[test]
96    fn estimate_cost_end_to_end_with_seeded_store() {
97        let dir = std::env::temp_dir().join(format!("obol-api-{}", std::process::id()));
98        std::env::set_var("OBOL_PRICING_DIR", &dir);
99        // seed the store from the sample sheet
100        let store = pricing::refresh::normalize_litellm(
101            include_bytes!("../tests/fixtures/litellm-sample.json"),
102            "2026-06-04",
103        )
104        .unwrap();
105        store.save(&pricing::current_path()).unwrap();
106
107        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
108        std::fs::write(
109            &tmp,
110            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
111        )
112        .unwrap();
113        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
114        assert!(est.total_usd > 0.0);
115        assert_eq!(est.pricing_as_of, "2026-06-04");
116        std::fs::remove_file(&tmp).ok();
117
118        std::fs::remove_dir_all(&dir).ok();
119        std::env::remove_var("OBOL_PRICING_DIR");
120    }
121
122    #[test]
123    fn estimate_cost_from_path_then_detect() {
124        let dir = std::env::temp_dir().join(format!("obol-path-{}", std::process::id()));
125        std::env::set_var("OBOL_PRICING_DIR", &dir);
126        let store = pricing::refresh::normalize_litellm(
127            include_bytes!("../tests/fixtures/litellm-sample.json"),
128            "2026-06-04",
129        )
130        .unwrap();
131        store.save(&pricing::current_path()).unwrap();
132
133        // Write the Claude fixture to a real file, detect dialect, then price.
134        let transcript = dir.join("session.jsonl");
135        std::fs::write(
136            &transcript,
137            include_bytes!("../tests/fixtures/claude-mini.jsonl"),
138        )
139        .unwrap();
140        let bytes = std::fs::read(&transcript).unwrap();
141        let d = transcript::detect(&bytes).unwrap();
142        let est = estimate_cost(&transcript, d).unwrap();
143        assert!(est.total_usd > 0.0);
144
145        std::fs::remove_dir_all(&dir).ok();
146        std::env::remove_var("OBOL_PRICING_DIR");
147    }
148
149    #[test]
150    fn falls_back_to_embedded_when_no_local_snapshot() {
151        // Force "no on-disk snapshot" hermetically: point XDG at an empty dir so
152        // `current_path()` resolves to a nonexistent file and the embedded snapshot
153        // is used. (Setting OBOL_PRICING_DIR instead would take the explicit-override
154        // branch and error PricingTablesMissing rather than fall back to embedded.)
155        let xdg = std::env::temp_dir().join(format!("obol-xdg-{}", std::process::id()));
156        std::fs::create_dir_all(&xdg).unwrap();
157        std::env::remove_var("OBOL_PRICING_DIR");
158        std::env::set_var("XDG_DATA_HOME", &xdg);
159        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
160        std::fs::write(
161            &tmp,
162            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
163        )
164        .unwrap();
165        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
166        assert_eq!(est.pricing_source, crate::model::PricingSource::Bundled);
167        assert!(est.total_usd > 0.0, "embedded snapshot should price claude");
168        std::fs::remove_file(&tmp).ok();
169        std::env::remove_var("XDG_DATA_HOME");
170        std::fs::remove_dir_all(&xdg).ok();
171    }
172
173    #[test]
174    fn explicit_override_uses_local_source() {
175        let dir = std::env::temp_dir().join(format!("obol-resolve-{}", std::process::id()));
176        std::env::set_var("OBOL_PRICING_DIR", &dir);
177        let store = pricing::refresh::normalize_litellm(
178            include_bytes!("../tests/fixtures/litellm-sample.json"),
179            "2099-01-01",
180        )
181        .unwrap();
182        store.save(&pricing::current_path()).unwrap();
183        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
184        std::fs::write(
185            &tmp,
186            include_bytes!("../tests/fixtures/claude-mini.jsonl").as_slice(),
187        )
188        .unwrap();
189        let est = estimate_cost(&tmp, Dialect::Claude).unwrap();
190        assert_eq!(est.pricing_source, crate::model::PricingSource::Local);
191        std::fs::remove_file(&tmp).ok();
192        std::fs::remove_dir_all(&dir).ok();
193        std::env::remove_var("OBOL_PRICING_DIR");
194    }
195
196    #[test]
197    fn kimi_model_surfaces_unpriced_loudly() {
198        std::env::remove_var("OBOL_PRICING_DIR");
199        let tmp = std::env::temp_dir().join(format!("obol-t-{}-{}", std::process::id(), line!()));
200        std::fs::write(
201            &tmp,
202            include_bytes!("../tests/fixtures/kimi-mini.jsonl").as_slice(),
203        )
204        .unwrap();
205        let est = estimate_cost(&tmp, Dialect::Kimi).unwrap();
206        assert_eq!(est.total_usd, 0.0, "kimi-for-coding is unpriced -> $0");
207        assert!(
208            est.unpriced_models.contains(&"kimi-for-coding".to_string()),
209            "must name the unpriced model: {:?}",
210            est.unpriced_models
211        );
212        std::fs::remove_file(&tmp).ok();
213    }
214
215    #[test]
216    fn refresh_report_serializes() {
217        let r = RefreshReport {
218            models: 7,
219            as_of: "2026-06-05".into(),
220            written_to: "/x/current.json".into(),
221        };
222        let v = serde_json::to_value(&r).unwrap();
223        assert_eq!(v["models"], 7);
224        assert_eq!(v["as_of"], "2026-06-05");
225        assert_eq!(v["written_to"], "/x/current.json");
226    }
227}
obol_core/lib.rs

obol_core/
lib.rs
