Skip to main content

cc_token_usage/
config.rs

1use crate::pricing::calculator::ModelPrice;
2use serde::Deserialize;
3use std::collections::HashMap;
4use std::path::Path;
5
6#[derive(Debug, Deserialize, Default)]
7pub struct Config {
8    #[serde(default)]
9    pub subscription: Vec<SubscriptionPeriod>,
10    #[serde(default)]
11    pub pricing_override: HashMap<String, PriceOverride>,
12}
13
14#[derive(Debug, Deserialize)]
15pub struct SubscriptionPeriod {
16    pub start_date: String,
17    pub monthly_price_usd: f64,
18    pub plan: Option<String>,
19}
20
21#[derive(Debug, Deserialize)]
22pub struct PriceOverride {
23    pub base_input: f64,
24    pub cache_write_5m: f64,
25    pub cache_write_1h: f64,
26    pub cache_read: f64,
27    pub output: f64,
28}
29
30impl Config {
31    pub fn load(path: &Path) -> anyhow::Result<Self> {
32        let content = std::fs::read_to_string(path)?;
33        Ok(toml::from_str(&content)?)
34    }
35
36    pub fn to_model_prices(&self) -> HashMap<String, ModelPrice> {
37        self.pricing_override
38            .iter()
39            .map(|(k, v)| {
40                (k.clone(), ModelPrice {
41                    base_input: v.base_input,
42                    cache_write_5m: v.cache_write_5m,
43                    cache_write_1h: v.cache_write_1h,
44                    cache_read: v.cache_read,
45                    output: v.output,
46                })
47            })
48            .collect()
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    #[test]
57    fn parse_config() {
58        let toml_str = r#"
59[[subscription]]
60start_date = "2026-01-01"
61monthly_price_usd = 100.0
62plan = "max_5x"
63
64[[subscription]]
65start_date = "2026-03-01"
66monthly_price_usd = 200.0
67plan = "max_20x"
68
69[pricing_override.claude-opus-4-6]
70base_input = 5.0
71cache_write_5m = 6.25
72cache_write_1h = 10.0
73cache_read = 0.50
74output = 25.0
75"#;
76
77        let config: Config = toml::from_str(toml_str).unwrap();
78
79        assert_eq!(config.subscription.len(), 2);
80        assert_eq!(config.subscription[0].start_date, "2026-01-01");
81        assert!((config.subscription[0].monthly_price_usd - 100.0).abs() < f64::EPSILON);
82        assert_eq!(config.subscription[0].plan.as_deref(), Some("max_5x"));
83        assert_eq!(config.subscription[1].start_date, "2026-03-01");
84        assert!((config.subscription[1].monthly_price_usd - 200.0).abs() < f64::EPSILON);
85        assert_eq!(config.subscription[1].plan.as_deref(), Some("max_20x"));
86
87        assert_eq!(config.pricing_override.len(), 1);
88        let opus = &config.pricing_override["claude-opus-4-6"];
89        assert!((opus.base_input - 5.0).abs() < f64::EPSILON);
90        assert!((opus.cache_write_5m - 6.25).abs() < f64::EPSILON);
91        assert!((opus.cache_write_1h - 10.0).abs() < f64::EPSILON);
92        assert!((opus.cache_read - 0.50).abs() < f64::EPSILON);
93        assert!((opus.output - 25.0).abs() < f64::EPSILON);
94
95        let model_prices = config.to_model_prices();
96        assert_eq!(model_prices.len(), 1);
97        let mp = &model_prices["claude-opus-4-6"];
98        assert!((mp.base_input - 5.0).abs() < f64::EPSILON);
99        assert!((mp.output - 25.0).abs() < f64::EPSILON);
100    }
101
102    #[test]
103    fn empty_config() {
104        let config: Config = toml::from_str("").unwrap();
105        assert!(config.subscription.is_empty());
106        assert!(config.pricing_override.is_empty());
107        assert!(config.to_model_prices().is_empty());
108    }
109}