use std::time::{Duration, Instant};
use kcode_codex_terra_usage::UsageAccumulator;
use rust_decimal::Decimal;
use serde_json::{Value, json};
fn counts(input: u64, cached: u64, cache_write: u64, output: u64, reasoning: u64) -> Value {
json!({
"inputTokens": input,
"cachedInputTokens": cached,
"cacheWriteInputTokens": cache_write,
"outputTokens": output,
"reasoningOutputTokens": reasoning
})
}
fn update(total: Value, last: Value) -> Value {
json!({"total": total, "last": last})
}
fn error_message(value: Value) -> String {
UsageAccumulator::new()
.apply(&value)
.unwrap_err()
.to_string()
}
fn assert_value(accumulator: &UsageAccumulator, key: &str, units: Decimal, price_cents: Decimal) {
let value = &accumulator.snapshot()[key];
assert_eq!(value.units, units);
assert_eq!(value.price_cents, price_cents);
}
#[test]
fn parse_errors_preserve_diagnostics() {
assert_eq!(
error_message(json!({"last": counts(0, 0, 0, 0, 0)})),
"Codex usage omitted cumulative totals"
);
assert_eq!(
error_message(json!({"total": counts(0, 0, 0, 0, 0)})),
"Codex usage omitted latest-round totals"
);
for key in [
"inputTokens",
"cachedInputTokens",
"cacheWriteInputTokens",
"outputTokens",
"reasoningOutputTokens",
] {
let mut missing = update(counts(1, 0, 0, 0, 0), counts(1, 0, 0, 0, 0));
missing["total"].as_object_mut().unwrap().remove(key);
assert_eq!(
error_message(missing),
format!("Codex usage contained invalid {key}")
);
let mut invalid = update(counts(1, 0, 0, 0, 0), counts(1, 0, 0, 0, 0));
invalid["last"]
.as_object_mut()
.unwrap()
.insert(key.to_owned(), json!("invalid"));
assert_eq!(
error_message(invalid),
format!("Codex usage contained invalid {key}")
);
}
}
#[test]
fn initial_and_duplicate_reconciliation_is_exact() {
let mut accumulator = UsageAccumulator::new();
let first = update(counts(10, 1, 0, 4, 2), counts(10, 1, 0, 4, 2));
accumulator.apply(&first).unwrap();
accumulator.apply(&first).unwrap();
assert_eq!(accumulator.reconciled_rounds(), 1);
let conflicting = update(counts(10, 1, 0, 4, 2), counts(0, 0, 0, 0, 0));
assert_eq!(
accumulator.apply(&conflicting).unwrap_err().to_string(),
"Codex duplicate usage changed its latest round"
);
assert_eq!(
error_message(update(counts(2, 0, 0, 0, 0), counts(1, 0, 0, 0, 0))),
"Codex initial usage did not reconcile"
);
}
#[test]
fn deltas_must_be_monotonic_and_exact() {
let mut accumulator = UsageAccumulator::new();
accumulator
.apply(&update(counts(10, 0, 0, 0, 0), counts(10, 0, 0, 0, 0)))
.unwrap();
assert_eq!(
accumulator
.apply(&update(counts(9, 0, 0, 0, 0), counts(0, 0, 0, 0, 0),))
.unwrap_err()
.to_string(),
"Codex usage delta did not reconcile"
);
let mut accumulator = UsageAccumulator::new();
accumulator
.apply(&update(counts(10, 0, 0, 0, 0), counts(10, 0, 0, 0, 0)))
.unwrap();
assert_eq!(
accumulator
.apply(&update(counts(15, 0, 0, 0, 0), counts(4, 0, 0, 0, 0),))
.unwrap_err()
.to_string(),
"Codex usage delta did not reconcile"
);
}
#[test]
fn impossible_totals_are_rejected() {
assert_eq!(
error_message(update(counts(1, 1, 1, 0, 0), counts(1, 1, 1, 0, 0),)),
"Codex usage contained impossible token totals"
);
assert_eq!(
error_message(update(counts(1, 0, 0, 1, 2), counts(1, 0, 0, 1, 2),)),
"Codex usage contained impossible token totals"
);
assert_eq!(
error_message(update(
counts(u64::MAX, u64::MAX, 1, 0, 0),
counts(u64::MAX, u64::MAX, 1, 0, 0),
)),
"Codex usage contained impossible token totals"
);
}
#[test]
fn zero_values_create_all_four_entries() {
let mut accumulator = UsageAccumulator::new();
accumulator
.apply(&update(counts(0, 0, 0, 0, 0), counts(0, 0, 0, 0, 0)))
.unwrap();
assert_eq!(accumulator.reconciled_rounds(), 1);
let snapshot = accumulator.snapshot();
assert_eq!(snapshot.len(), 4);
for value in snapshot.values() {
assert_eq!(value.units, Decimal::ZERO);
assert_eq!(value.price_cents, Decimal::ZERO);
}
}
#[test]
fn cache_writes_and_tier_boundaries_use_exact_decimals() {
let mut cache_writes = UsageAccumulator::new();
cache_writes
.apply(&update(counts(10, 3, 2, 9, 4), counts(10, 3, 2, 9, 4)))
.unwrap();
assert_value(
&cache_writes,
"input tokens",
Decimal::from(7),
Decimal::new(15, 4),
);
assert_value(
&cache_writes,
"cached input tokens",
Decimal::from(3),
Decimal::new(6, 5),
);
assert_value(
&cache_writes,
"output tokens",
Decimal::from(5),
Decimal::new(6, 3),
);
assert_value(
&cache_writes,
"reasoning tokens",
Decimal::from(4),
Decimal::new(48, 4),
);
let mut low = UsageAccumulator::new();
low.apply(&update(
counts(272_000, 0, 0, 0, 0),
counts(272_000, 0, 0, 0, 0),
))
.unwrap();
assert_value(
&low,
"input tokens",
Decimal::from(272_000),
Decimal::new(544, 1),
);
let mut high = UsageAccumulator::new();
high.apply(&update(
counts(272_001, 0, 0, 0, 0),
counts(272_001, 0, 0, 0, 0),
))
.unwrap();
assert_value(
&high,
"input tokens",
Decimal::from(272_001),
Decimal::new(1_088_004, 4),
);
}
#[test]
fn multiple_rounds_accumulate_and_snapshot_is_owned() {
let mut accumulator = UsageAccumulator::new();
accumulator
.apply(&update(
counts(100, 20, 10, 40, 10),
counts(100, 20, 10, 40, 10),
))
.unwrap();
accumulator
.apply(&update(
counts(272_101, 20, 11, 44, 11),
counts(272_001, 0, 1, 4, 1),
))
.unwrap();
assert_eq!(accumulator.reconciled_rounds(), 2);
assert_value(
&accumulator,
"input tokens",
Decimal::from(272_081),
Decimal::new(1_088_170, 4),
);
assert_value(
&accumulator,
"cached input tokens",
Decimal::from(20),
Decimal::new(4, 4),
);
assert_value(
&accumulator,
"output tokens",
Decimal::from(33),
Decimal::new(414, 4),
);
assert_value(
&accumulator,
"reasoning tokens",
Decimal::from(11),
Decimal::new(138, 4),
);
let mut snapshot = accumulator.snapshot();
snapshot.clear();
assert_eq!(accumulator.snapshot().len(), 4);
}
#[test]
fn accumulator_canary() {
let start = Instant::now();
let mut accumulator = UsageAccumulator::new();
for total in 1_u64..=100_000 {
accumulator
.apply(&update(counts(total, 0, 0, 0, 0), counts(1, 0, 0, 0, 0)))
.unwrap();
}
let limit = if cfg!(debug_assertions) {
Duration::from_secs(5)
} else {
Duration::from_secs(1)
};
assert!(start.elapsed() < limit);
assert_eq!(accumulator.reconciled_rounds(), 100_000);
}