kcode_codex_terra_usage/
lib.rs

1use std::{collections::BTreeMap, fmt};
2
3use kcode_k1_accounting::UsageValue;
4use rust_decimal::Decimal;
5use serde_json::Value;
6
7type RoundUsage = [u64; 5];
8
9#[derive(Clone, Debug, Eq, PartialEq)]
10pub struct Error {
11    message: String,
12}
13
14impl Error {
15    fn new(message: impl Into<String>) -> Self {
16        Self {
17            message: message.into(),
18        }
19    }
20}
21
22impl fmt::Display for Error {
23    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
24        formatter.write_str(&self.message)
25    }
26}
27
28impl std::error::Error for Error {}
29
30#[derive(Clone, Debug)]
31pub struct UsageAccumulator {
32    usage: BTreeMap<String, UsageValue>,
33    total: Option<RoundUsage>,
34    last: Option<RoundUsage>,
35    reconciled_rounds: usize,
36}
37
38impl UsageAccumulator {
39    pub fn new() -> Self {
40        Self {
41            usage: BTreeMap::new(),
42            total: None,
43            last: None,
44            reconciled_rounds: 0,
45        }
46    }
47
48    pub fn apply(&mut self, value: &Value) -> Result<(), Error> {
49        let total = parse(
50            value
51                .get("total")
52                .ok_or_else(|| Error::new("Codex usage omitted cumulative totals"))?,
53        )?;
54        let last = parse(
55            value
56                .get("last")
57                .ok_or_else(|| Error::new("Codex usage omitted latest-round totals"))?,
58        )?;
59
60        if let Some(previous) = self.total {
61            if total == previous {
62                if self.last != Some(last) {
63                    return Err(Error::new("Codex duplicate usage changed its latest round"));
64                }
65                return Ok(());
66            }
67            if delta(total, previous) != Some(last) {
68                return Err(Error::new("Codex usage delta did not reconcile"));
69            }
70        } else if total != last {
71            return Err(Error::new("Codex initial usage did not reconcile"));
72        }
73
74        self.add(last);
75        self.total = Some(total);
76        self.last = Some(last);
77        self.reconciled_rounds = self.reconciled_rounds.saturating_add(1);
78        Ok(())
79    }
80
81    pub fn reconciled_rounds(&self) -> usize {
82        self.reconciled_rounds
83    }
84
85    pub fn snapshot(&self) -> BTreeMap<String, UsageValue> {
86        self.usage.clone()
87    }
88
89    fn add(&mut self, round: RoundUsage) {
90        let [input, cached, cache_write, output, reasoning] = round;
91        let ordinary = input - cached - cache_write;
92        let (ordinary_rate, cache_write_rate, cached_rate, output_rate) = if input > 272_000 {
93            (400, 500, 40, 1_800)
94        } else {
95            (200, 250, 20, 1_200)
96        };
97
98        for (key, units, rate) in [
99            ("input tokens", ordinary, ordinary_rate),
100            ("input tokens", cache_write, cache_write_rate),
101            ("cached input tokens", cached, cached_rate),
102            ("output tokens", output - reasoning, output_rate),
103            ("reasoning tokens", reasoning, output_rate),
104        ] {
105            add(&mut self.usage, key, units, rate);
106        }
107    }
108}
109
110impl Default for UsageAccumulator {
111    fn default() -> Self {
112        Self::new()
113    }
114}
115
116fn delta(current: RoundUsage, previous: RoundUsage) -> Option<RoundUsage> {
117    Some([
118        current[0].checked_sub(previous[0])?,
119        current[1].checked_sub(previous[1])?,
120        current[2].checked_sub(previous[2])?,
121        current[3].checked_sub(previous[3])?,
122        current[4].checked_sub(previous[4])?,
123    ])
124}
125
126fn parse(value: &Value) -> Result<RoundUsage, Error> {
127    let part = [
128        integer(value, "inputTokens")?,
129        integer(value, "cachedInputTokens")?,
130        integer(value, "cacheWriteInputTokens")?,
131        integer(value, "outputTokens")?,
132        integer(value, "reasoningOutputTokens")?,
133    ];
134
135    if part[1]
136        .checked_add(part[2])
137        .is_none_or(|value| value > part[0])
138        || part[4] > part[3]
139    {
140        return Err(Error::new("Codex usage contained impossible token totals"));
141    }
142
143    Ok(part)
144}
145
146fn integer(value: &Value, key: &str) -> Result<u64, Error> {
147    value
148        .get(key)
149        .and_then(Value::as_u64)
150        .ok_or_else(|| Error::new(format!("Codex usage contained invalid {key}")))
151}
152
153fn add(usage: &mut BTreeMap<String, UsageValue>, key: &str, units: u64, rate: u64) {
154    let entry = usage.entry(key.to_owned()).or_default();
155    let units = Decimal::from(units);
156    entry.units += units;
157    entry.price_cents += units * Decimal::from(rate) / Decimal::from(1_000_000_u64);
158}
kcode_codex_terra_usage/lib.rs

kcode_codex_terra_usage/
lib.rs
