use std::cell::RefCell;
use std::collections::HashMap;
use std::time::{Duration, Instant};
use super::cache::{
CacheReadError, CacheWriteError, load_raw_cache_if_fresh, load_raw_cache_snapshot,
save_raw_cache,
};
use super::provider::fetch_litellm_raw;
use super::resolver::{fallback_pricing, parse_litellm_data, resolve_pricing_known};
use super::source::{CacheMetadata, PricingSource};
use super::types::ModelPricing;
#[derive(Debug, Clone)]
enum ResolvedPricing {
Known {
pricing: ModelPricing,
source: PricingSource,
},
Unknown,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum PricingLoadError {
#[error("failed to load pricing cache: {0}")]
Cache(#[from] CacheReadError),
}
#[derive(Debug)]
pub(crate) struct PricingDb {
models: HashMap<String, ModelPricing>,
resolved: RefCell<HashMap<String, ResolvedPricing>>,
strict_unknown: bool,
source: PricingSource,
cache_metadata: Option<CacheMetadata>,
}
const PRICING_CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
impl PricingDb {
fn empty(strict_unknown: bool) -> Self {
Self {
models: HashMap::new(),
resolved: RefCell::new(HashMap::new()),
strict_unknown,
source: PricingSource::Fallback,
cache_metadata: None,
}
}
fn from_raw_data(
data: HashMap<String, serde_json::Value>,
strict_unknown: bool,
source: PricingSource,
cache_metadata: Option<CacheMetadata>,
) -> Self {
Self {
models: parse_litellm_data(data),
resolved: RefCell::new(HashMap::new()),
strict_unknown,
source,
cache_metadata,
}
}
fn load_from_cache(strict_unknown: bool) -> Result<Option<Self>, CacheReadError> {
Ok(load_raw_cache_snapshot()?.map(|snapshot| {
let source = if snapshot.metadata.age > PRICING_CACHE_TTL {
PricingSource::CacheStale
} else {
PricingSource::Cache
};
Self::from_raw_data(
snapshot.data,
strict_unknown,
source,
Some(snapshot.metadata),
)
}))
}
fn load_from_cache_if_fresh(
ttl: Duration,
strict_unknown: bool,
) -> Result<Option<(Self, Duration)>, CacheReadError> {
Ok(load_raw_cache_if_fresh(ttl)?.map(|snapshot| {
(
Self::from_raw_data(
snapshot.data,
strict_unknown,
PricingSource::Cache,
Some(snapshot.metadata),
),
snapshot.metadata.age,
)
}))
}
pub(crate) fn load(offline: bool, strict_unknown: bool) -> Self {
Self::try_load(offline, strict_unknown).unwrap_or_else(|error| {
eprintln!("Error: {error}");
std::process::exit(1);
})
}
pub(crate) fn load_quiet(offline: bool, strict_unknown: bool) -> Self {
Self::try_load_quiet(offline, strict_unknown).unwrap_or_else(|error| {
eprintln!("Error: {error}");
std::process::exit(1);
})
}
pub(crate) fn try_load(offline: bool, strict_unknown: bool) -> Result<Self, PricingLoadError> {
Self::load_internal(offline, strict_unknown, false)
}
pub(crate) fn try_load_quiet(
offline: bool,
strict_unknown: bool,
) -> Result<Self, PricingLoadError> {
Self::load_internal(offline, strict_unknown, true)
}
fn load_internal(
offline: bool,
strict_unknown: bool,
quiet: bool,
) -> Result<Self, PricingLoadError> {
let start = Instant::now();
if offline {
return Self::finish_offline_cache_load(
Self::load_from_cache(strict_unknown),
strict_unknown,
quiet,
start,
);
}
match Self::load_from_cache_if_fresh(PRICING_CACHE_TTL, strict_unknown) {
Ok(Some((db, age))) => {
if !quiet {
eprintln!(
"Using cached pricing ({:.1}h old)",
age.as_secs_f64() / 3600.0
);
}
return Ok(db);
}
Ok(None) => {}
Err(error) => {
eprintln!("Warning: ignoring invalid pricing cache before refresh: {error}");
}
}
if !quiet {
eprint!("Fetching pricing from LiteLLM...");
}
if let Some(raw_data) = fetch_litellm_raw() {
let fetch_time = start.elapsed();
let save_result = save_raw_cache(&raw_data);
let db = Self::from_raw_data(raw_data, strict_unknown, PricingSource::Live, None);
if !quiet {
eprintln!(
" {} models ({:.2}ms)",
db.models.len(),
fetch_time.as_secs_f64() * 1000.0
);
}
warn_cache_write_error(save_result);
return Ok(db);
}
if !quiet {
eprintln!(" failed, trying cache...");
}
match Self::load_from_cache(strict_unknown) {
Ok(Some(db)) => {
if !quiet {
eprintln!(
"Using cached pricing ({:.2}ms)",
start.elapsed().as_secs_f64() * 1000.0
);
}
return Ok(db);
}
Ok(None) => {}
Err(error) => return Err(error.into()),
}
if !quiet {
eprintln!(
"Using defaults ({:.2}ms)",
start.elapsed().as_secs_f64() * 1000.0
);
}
Ok(Self::empty(strict_unknown))
}
fn finish_offline_cache_load(
cache_result: Result<Option<Self>, CacheReadError>,
strict_unknown: bool,
quiet: bool,
start: Instant,
) -> Result<Self, PricingLoadError> {
match cache_result {
Ok(Some(db)) => {
if !quiet {
eprintln!(
"Using cached pricing ({:.2}ms)",
start.elapsed().as_secs_f64() * 1000.0
);
}
Ok(db)
}
Ok(None) => {
if !quiet {
eprintln!(
"No cached pricing, using defaults ({:.2}ms)",
start.elapsed().as_secs_f64() * 1000.0
);
}
Ok(Self::empty(strict_unknown))
}
Err(error) => Err(error.into()),
}
}
pub(super) fn get_pricing(&self, model: &str) -> Option<ModelPricing> {
self.resolve_pricing(model).map(|(pricing, _)| pricing)
}
pub(crate) fn pricing_source_for_model(&self, model: &str) -> Option<PricingSource> {
self.resolve_pricing(model).map(|(_, source)| source)
}
pub(crate) fn source(&self) -> PricingSource {
self.source
}
pub(crate) fn cache_age_seconds(&self) -> Option<u64> {
self.cache_metadata.map(CacheMetadata::age_seconds)
}
pub(crate) fn cache_modified_epoch_seconds(&self) -> Option<u64> {
self.cache_metadata
.map(CacheMetadata::modified_epoch_seconds)
}
#[cfg(test)]
pub(super) fn insert_model_for_tests(&mut self, name: String, pricing: ModelPricing) {
self.models.insert(name, pricing);
}
fn resolve_pricing(&self, model: &str) -> Option<(ModelPricing, PricingSource)> {
if let Some(cached) = self.resolved.borrow().get(model) {
return match cached {
ResolvedPricing::Known { pricing, source } => Some((pricing.clone(), *source)),
ResolvedPricing::Unknown => None,
};
}
let pricing = if let Some(pricing) = resolve_pricing_known(model, &self.models) {
Some((pricing, self.source))
} else if self.strict_unknown {
None
} else {
fallback_pricing(model).map(|pricing| (pricing, PricingSource::Fallback))
};
let cached = match &pricing {
Some((pricing, source)) => ResolvedPricing::Known {
pricing: pricing.clone(),
source: *source,
},
None => ResolvedPricing::Unknown,
};
self.resolved.borrow_mut().insert(model.to_string(), cached);
pricing
}
}
fn warn_cache_write_error(result: Result<(), CacheWriteError>) {
if let Err(error) = result {
eprintln!("Warning: failed to save pricing cache: {error}");
}
}
impl Default for PricingDb {
fn default() -> Self {
Self {
models: HashMap::new(),
resolved: RefCell::new(HashMap::new()),
strict_unknown: false,
source: PricingSource::Fallback,
cache_metadata: None,
}
}
}
#[cfg(test)]
#[allow(clippy::float_cmp)]
mod tests {
use super::*;
use crate::core::Stats;
use crate::pricing::{
attach_costs, calculate_cost, pricing_source_for_models, sum_model_costs,
};
use serde_json::json;
use std::fs;
use std::path::PathBuf;
use std::time::SystemTime;
use tempfile::TempDir;
fn sample_raw_pricing(model: &str) -> HashMap<String, serde_json::Value> {
HashMap::from([(
model.to_string(),
json!({
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
}),
)])
}
fn malformed_cache_error() -> CacheReadError {
let source =
serde_json::from_str::<HashMap<String, serde_json::Value>>("{not json").unwrap_err();
CacheReadError::Malformed {
path: PathBuf::from("pricing.json"),
source,
}
}
#[test]
fn offline_missing_cache_keeps_default_pricing_behavior() {
let db = PricingDb::finish_offline_cache_load(Ok(None), true, true, Instant::now())
.expect("missing cache should use defaults");
assert!(db.models.is_empty());
assert!(db.strict_unknown);
}
#[test]
fn offline_malformed_cache_fails_closed() {
let error = PricingDb::finish_offline_cache_load(
Err(malformed_cache_error()),
false,
true,
Instant::now(),
)
.unwrap_err();
assert!(error.to_string().contains("malformed"));
}
#[test]
fn cache_read_distinguishes_missing_from_malformed_for_db_load() {
let root = TempDir::new().unwrap();
let missing_path = root.path().join("missing-pricing.json");
assert!(
super::super::cache::load_raw_cache_from_paths(&[missing_path])
.unwrap()
.is_none()
);
let malformed_path = root.path().join("pricing.json");
fs::write(&malformed_path, "{not json").unwrap();
let error = super::super::cache::load_raw_cache_from_paths(&[malformed_path]).unwrap_err();
assert!(matches!(error, CacheReadError::Malformed { .. }));
}
#[test]
fn fetched_pricing_remains_usable_after_cache_save_failure() {
let root = TempDir::new().unwrap();
let blocker = root.path().join("not-a-directory");
fs::write(&blocker, "file").unwrap();
let cache_path = blocker.join("pricing.json");
let raw_data = sample_raw_pricing("gpt-5");
let save_result = super::super::cache::save_raw_cache_to_path(&raw_data, &cache_path);
let db = PricingDb::from_raw_data(raw_data, false, PricingSource::Live, None);
assert!(save_result.is_err());
assert!(db.get_pricing("gpt-5").is_some());
}
#[test]
fn calculate_cost_basic() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
reasoning_output: 15e-6,
cache_create: 3.75e-6,
cache_create_1h: 3.75e-6,
cache_read: 0.3e-6,
},
);
let stats = Stats {
input_tokens: 1_000_000,
output_tokens: 100_000,
cache_creation: 0,
cache_creation_1h: 0,
cache_read: 0,
reasoning_tokens: 0,
count: 1,
skipped_chunks: 0,
estimated_proxy: crate::core::CostTokens::default(),
};
let cost = calculate_cost(&stats, "sonnet-4", &db);
assert!((cost - 4.5).abs() < 0.001);
}
#[test]
fn calculate_cost_with_cache() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
reasoning_output: 15e-6,
cache_create: 3.75e-6,
cache_create_1h: 3.75e-6,
cache_read: 0.3e-6,
},
);
let stats = Stats {
input_tokens: 0,
output_tokens: 0,
cache_creation: 1_000_000,
cache_creation_1h: 0,
cache_read: 1_000_000,
reasoning_tokens: 0,
count: 1,
skipped_chunks: 0,
estimated_proxy: crate::core::CostTokens::default(),
};
let cost = calculate_cost(&stats, "sonnet-4", &db);
assert!((cost - 4.05).abs() < 0.001);
}
#[test]
fn calculate_cost_zero_tokens() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
..Default::default()
},
);
let stats = Stats::default();
let cost = calculate_cost(&stats, "sonnet-4", &db);
assert_eq!(cost, 0.0);
}
#[test]
fn get_pricing_fallback_for_unknown_model() {
let db = PricingDb::default();
let pricing = db.get_pricing("sonnet-4");
assert!(pricing.as_ref().is_some_and(|p| p.input > 0.0));
assert!(pricing.as_ref().is_some_and(|p| p.output > 0.0));
}
#[test]
fn strict_mode_marks_unknown_model_as_nan_cost() {
let db = PricingDb {
strict_unknown: true,
..PricingDb::default()
};
let stats = Stats {
input_tokens: 10,
..Default::default()
};
let cost = calculate_cost(&stats, "totally-unknown-model", &db);
assert!(cost.is_nan());
}
#[test]
fn calculate_cost_with_reasoning_tokens() {
let mut db = PricingDb::default();
db.models.insert(
"opus-4".to_string(),
ModelPricing {
input: 15e-6,
output: 75e-6,
reasoning_output: 75e-6,
cache_create: 0.0,
cache_create_1h: 0.0,
cache_read: 0.0,
},
);
let stats = Stats {
input_tokens: 100_000,
output_tokens: 0,
reasoning_tokens: 50_000,
..Default::default()
};
let cost = calculate_cost(&stats, "opus-4", &db);
assert!((cost - 5.25).abs() < 0.001);
}
#[test]
fn sum_model_costs_multiple_models() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
..Default::default()
},
);
db.models.insert(
"haiku-3.5".to_string(),
ModelPricing {
input: 0.8e-6,
output: 4e-6,
..Default::default()
},
);
let mut models = HashMap::new();
models.insert(
"sonnet-4".to_string(),
Stats {
input_tokens: 1_000_000,
output_tokens: 100_000,
..Default::default()
},
);
models.insert(
"haiku-3.5".to_string(),
Stats {
input_tokens: 500_000,
output_tokens: 50_000,
..Default::default()
},
);
let total = sum_model_costs(&models, &db);
assert!((total - 5.1).abs() < 0.001);
}
#[test]
fn sum_model_costs_empty_map() {
let db = PricingDb::default();
let models: HashMap<String, Stats> = HashMap::new();
let total = sum_model_costs(&models, &db);
assert_eq!(total, 0.0);
}
#[test]
fn sum_model_costs_returns_nan_when_all_unknown() {
let db = PricingDb {
strict_unknown: true,
..PricingDb::default()
};
let mut models = HashMap::new();
models.insert(
"totally-unknown-xyz".to_string(),
Stats {
input_tokens: 100,
..Default::default()
},
);
let total = sum_model_costs(&models, &db);
assert!(total.is_nan());
}
#[test]
fn sum_model_costs_skips_unknown_keeps_known() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
..Default::default()
},
);
let mut models = HashMap::new();
models.insert(
"sonnet-4".to_string(),
Stats {
input_tokens: 1_000_000,
output_tokens: 0,
..Default::default()
},
);
models.insert(
"totally-unknown-xyz".to_string(),
Stats {
input_tokens: 999_999_999, ..Default::default()
},
);
let total = sum_model_costs(&models, &db);
assert!((total - 3.0).abs() < 0.001);
}
#[test]
fn attach_costs_computes_per_item() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
..Default::default()
},
);
let items: Vec<(String, HashMap<String, Stats>)> = vec![
(
"day1".to_string(),
HashMap::from([(
"sonnet-4".to_string(),
Stats {
input_tokens: 1_000_000,
output_tokens: 0,
..Default::default()
},
)]),
),
(
"day2".to_string(),
HashMap::from([(
"sonnet-4".to_string(),
Stats {
input_tokens: 0,
output_tokens: 100_000,
..Default::default()
},
)]),
),
];
let costed = attach_costs(&items, |item| &item.1, &db);
assert_eq!(costed.len(), 2);
assert!((costed[0].cost - 3.0).abs() < 0.001);
assert!((costed[1].cost - 1.5).abs() < 0.001);
assert_eq!(costed[0].item.0, "day1");
assert_eq!(costed[1].item.0, "day2");
}
#[test]
fn attach_costs_empty_slice() {
let db = PricingDb::default();
let items: Vec<(String, HashMap<String, Stats>)> = vec![];
let costed = attach_costs(&items, |item| &item.1, &db);
assert!(costed.is_empty());
}
#[test]
fn get_pricing_caches_resolved_result() {
let mut db = PricingDb::default();
db.models.insert(
"sonnet-4".to_string(),
ModelPricing {
input: 3e-6,
output: 15e-6,
..Default::default()
},
);
let p1 = db.get_pricing("sonnet-4");
assert!(p1.is_some());
assert!(db.resolved.borrow().contains_key("sonnet-4"));
let p2 = db.get_pricing("sonnet-4");
assert_eq!(p1.unwrap().input, p2.unwrap().input);
}
#[test]
fn pricing_source_tracks_live_known_model() {
let db = PricingDb::from_raw_data(
sample_raw_pricing("claude-3-5-sonnet-20241022"),
false,
PricingSource::Live,
None,
);
assert_eq!(
db.pricing_source_for_model("claude-3-5-sonnet-20241022"),
Some(PricingSource::Live)
);
}
#[test]
fn pricing_source_tracks_stale_cache_known_model() {
let metadata = CacheMetadata {
age: Duration::from_secs(2 * 24 * 60 * 60),
modified: SystemTime::now() - Duration::from_secs(2 * 24 * 60 * 60),
};
let db = PricingDb::from_raw_data(
sample_raw_pricing("claude-3-5-sonnet-20241022"),
false,
PricingSource::CacheStale,
Some(metadata),
);
assert_eq!(
db.pricing_source_for_model("claude-3-5-sonnet-20241022"),
Some(PricingSource::CacheStale)
);
assert!(matches!(db.cache_age_seconds(), Some(age) if age >= 24 * 60 * 60));
}
#[test]
fn pricing_source_combines_known_and_fallback_as_mixed() {
let db = PricingDb::from_raw_data(
sample_raw_pricing("claude-3-5-sonnet-20241022"),
false,
PricingSource::Cache,
None,
);
let models = HashMap::from([
(
"claude-3-5-sonnet-20241022".to_string(),
Stats {
input_tokens: 100,
..Default::default()
},
),
(
"gpt-5".to_string(),
Stats {
input_tokens: 100,
..Default::default()
},
),
]);
assert_eq!(
pricing_source_for_models(&models, &db),
PricingSource::Mixed
);
}
#[test]
fn get_pricing_caches_unknown_in_strict_mode() {
let db = PricingDb {
strict_unknown: true,
..PricingDb::default()
};
let p = db.get_pricing("nonexistent-xyz-model");
assert!(p.is_none());
let resolved = db.resolved.borrow();
assert!(resolved.contains_key("nonexistent-xyz-model"));
assert!(matches!(
resolved.get("nonexistent-xyz-model"),
Some(ResolvedPricing::Unknown)
));
}
#[test]
fn default_pricing_db_has_empty_models() {
let db = PricingDb::default();
assert!(db.models.is_empty());
assert!(db.resolved.borrow().is_empty());
assert!(!db.strict_unknown);
}
}