kcode-codex-terra 0.1.0

One ephemeral GPT-5.6 Terra dynamic-tool inference with mandatory accounting.
Documentation
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);
}