use std::collections::BTreeMap;
use kcode_k1_accounting::{Accounting, AccountingEvent, UsageValue};
use rust_decimal::Decimal;
use serde_json::Value;
use crate::{Result, protocol};
type RoundUsage = [u64; 5];
fn delta(current: RoundUsage, previous: RoundUsage) -> Option<RoundUsage> {
Some([
current[0].checked_sub(previous[0])?,
current[1].checked_sub(previous[1])?,
current[2].checked_sub(previous[2])?,
current[3].checked_sub(previous[3])?,
current[4].checked_sub(previous[4])?,
])
}
pub(crate) struct AttemptAccounting {
accounting: Accounting,
usage: BTreeMap<String, UsageValue>,
total: Option<RoundUsage>,
last: Option<RoundUsage>,
reconciled_rounds: usize,
}
impl AttemptAccounting {
pub(crate) fn new(accounting: Accounting) -> Self {
Self {
accounting,
usage: BTreeMap::new(),
total: None,
last: None,
reconciled_rounds: 0,
}
}
pub(crate) fn apply(&mut self, value: &Value) -> Result<()> {
let total = parse(
value
.get("total")
.ok_or_else(|| protocol("Codex usage omitted cumulative totals"))?,
)?;
let last = parse(
value
.get("last")
.ok_or_else(|| protocol("Codex usage omitted latest-round totals"))?,
)?;
if let Some(previous) = self.total {
if total == previous {
if self.last != Some(last) {
return Err(protocol("Codex duplicate usage changed its latest round"));
}
return Ok(());
}
if delta(total, previous) != Some(last) {
return Err(protocol("Codex usage delta did not reconcile"));
}
} else if total != last {
return Err(protocol("Codex initial usage did not reconcile"));
}
self.add(last);
self.total = Some(total);
self.last = Some(last);
self.reconciled_rounds = self.reconciled_rounds.saturating_add(1);
Ok(())
}
pub(crate) fn reconciled_rounds(&self) -> usize {
self.reconciled_rounds
}
pub(crate) fn snapshot(&self) -> BTreeMap<String, UsageValue> {
self.usage.clone()
}
fn add(&mut self, round: RoundUsage) {
let [input, cached, cache_write, output, reasoning] = round;
let high = input > 272_000;
let ordinary = input - cached - cache_write;
let (ordinary_rate, cache_write_rate, cached_rate, output_rate) = if high {
(400, 500, 40, 1_800)
} else {
(200, 250, 20, 1_200)
};
for (key, units, rate) in [
("input tokens", ordinary, ordinary_rate),
("input tokens", cache_write, cache_write_rate),
("cached input tokens", cached, cached_rate),
("output tokens", output - reasoning, output_rate),
("reasoning tokens", reasoning, output_rate),
] {
add(&mut self.usage, key, units, rate);
}
}
}
impl Drop for AttemptAccounting {
fn drop(&mut self) {
self.accounting.record(&AccountingEvent {
source: super::MODEL.to_owned(),
operation: super::OPERATION.to_owned(),
usage: self.usage.clone(),
});
}
}
fn parse(value: &Value) -> Result<RoundUsage> {
let part = [
integer(value, "inputTokens")?,
integer(value, "cachedInputTokens")?,
integer(value, "cacheWriteInputTokens")?,
integer(value, "outputTokens")?,
integer(value, "reasoningOutputTokens")?,
];
if part[1]
.checked_add(part[2])
.is_none_or(|value| value > part[0])
|| part[4] > part[3]
{
return Err(protocol("Codex usage contained impossible token totals"));
}
Ok(part)
}
fn integer(value: &Value, key: &str) -> Result<u64> {
value
.get(key)
.and_then(Value::as_u64)
.ok_or_else(|| protocol(format!("Codex usage contained invalid {key}")))
}
fn add(usage: &mut BTreeMap<String, UsageValue>, key: &str, units: u64, rate: u64) {
let entry = usage.entry(key.to_owned()).or_default();
let units = Decimal::from(units);
entry.units += units;
entry.price_cents += units * Decimal::from(rate) / Decimal::from(1_000_000_u64);
}