basis 0.4.2

//! What a run reported spending.
//!
//! Its own module because it is the thing everything else about cost is built
//! on: the per-run bounds of ADR-0014 enforce a budget, and ADR-0010's shared
//! `BudgetPool` will draw from one — but both need somewhere honest to read
//! *what a turn actually cost*, and this is it.

/// What a run reported spending, summed over the rounds it took.
///
/// mentra emits one usage report per completed model response, carrying *that
/// round's* numbers rather than a running total — its own soft token budget is
/// a sum of them, and so is this. A turn that took four rounds reports four
/// times and lands here as one figure.
///
/// Three things to know before treating this as a bill.
///
/// It is *reported*, not measured. A provider that says nothing leaves this at
/// zero, and a stream that lagged — the [`Notice`](crate::Event::Notice) says
/// how many events were dropped — may have lost a report along with everything
/// else, which undercounts.
///
/// It counts the rounds of the run's own agent *and* of the work that agent
/// delegates. mentra's `task` intrinsic drives its subagent on the parent
/// run's [`RunOptions::child`](mentra::runtime::RunOptions::child) and relays
/// the child's usage reports onto the parent's bus, so a delegated round
/// arrives here like any other — indistinguishable from the parent's own,
/// which is consistent with the accounting being aggregate: mentra's usage
/// report carries no agent id for anyone to attribute by. Summing the stream
/// and reading this figure therefore give one answer, and it is the same
/// answer [`TurnOptions::token_budget`](super::TurnOptions) and
/// [`BudgetPool`](crate::BudgetPool) are enforced against, because the child
/// shares the parent's accounting handle rather than getting one of its own.
/// None of that was true before mentra `0436bae`: the child ran on fresh
/// options and reported to a bus nobody read, so a delegating run spent tokens
/// neither this figure nor any bound ever saw.
///
/// One edge survives, and it fails loudly rather than leaking: a delegation
/// issued once the shared budget is already crossed inherits an allowance with
/// nothing left in it, does zero rounds, and comes back as a failed tool call
/// instead of an empty success. That is the round-boundary softness
/// [`TurnOptions::token_budget`](super::TurnOptions) already describes, seen
/// from the delegating side, and mentra pins it
/// (`delegating_with_the_budget_already_spent_fails_the_delegation`).
///
/// And cache tokens are counted but never mixed in — see
/// [`total_tokens`](Self::total_tokens).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RunUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_creation_tokens: u64,
}

impl RunUsage {
    /// Input plus output — what [`RunConfig::with_token_budget`](super::RunConfig) counts.
    ///
    /// Cache reads and cache writes are left out for the same reason mentra
    /// leaves them out of the budget it enforces: they are priced differently
    /// everywhere, and a total that mixed them would answer no question
    /// exactly.
    pub const fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens
    }

    /// This total plus `other`'s — how a caller adds up several runs.
    ///
    /// The shape a shared budget wants: a fan-out folds its reports into one
    /// figure without any of the runs knowing about the others.
    #[must_use]
    pub const fn plus(self, other: Self) -> Self {
        Self {
            input_tokens: self.input_tokens + other.input_tokens,
            output_tokens: self.output_tokens + other.output_tokens,
            cache_read_tokens: self.cache_read_tokens + other.cache_read_tokens,
            cache_creation_tokens: self.cache_creation_tokens + other.cache_creation_tokens,
        }
    }

    /// This total plus whatever `event` reported, unchanged when it reported
    /// nothing.
    ///
    /// Read from the session event rather than from basis's mapped
    /// [`Event::Usage`](crate::Event::Usage) so that a sink which stopped
    /// accepting events still gets an honest total: what the run spent is not
    /// contingent on anyone having listened.
    pub(crate) fn recording(self, event: &mentra::SessionEvent) -> Self {
        let mentra::SessionEvent::UsageReport {
            input_tokens,
            output_tokens,
            cache_read_tokens,
            cache_creation_tokens,
            ..
        } = event
        else {
            return self;
        };

        Self {
            input_tokens: self.input_tokens + *input_tokens,
            output_tokens: self.output_tokens + *output_tokens,
            cache_read_tokens: self.cache_read_tokens + *cache_read_tokens,
            cache_creation_tokens: self.cache_creation_tokens + *cache_creation_tokens,
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;

    /// One round's worth of reported usage, as mentra emits it.
    fn usage_report(input: u64, output: u64) -> mentra::SessionEvent {
        mentra::SessionEvent::UsageReport {
            agent_id: "a1".to_string(),
            input_tokens: input,
            output_tokens: output,
            cache_read_tokens: 1,
            cache_creation_tokens: 2,
        }
    }

    #[test]
    fn usage_is_summed_over_the_rounds_of_a_turn() {
        // mentra reports per completed model response, not cumulatively — its
        // own token budget adds the reports up, and a basis that treated the
        // last one as the total would under-report every multi-round turn.
        let usage = RunUsage::default()
            .recording(&usage_report(100, 20))
            .recording(&usage_report(150, 30));

        assert_eq!(usage.input_tokens, 250);
        assert_eq!(usage.output_tokens, 50);
        assert_eq!(usage.cache_read_tokens, 2);
        assert_eq!(usage.cache_creation_tokens, 4);
        assert_eq!(usage.total_tokens(), 300, "the budget counts these two");
    }

    #[test]
    fn an_event_that_reports_no_usage_changes_nothing() {
        let counted = RunUsage::default().recording(&usage_report(10, 5));
        let after = counted.recording(&mentra::SessionEvent::UserMessage {
            text: "hello".to_string(),
        });

        assert_eq!(after, counted);
    }

    #[test]
    fn a_turn_that_reported_nothing_reports_zero() {
        // A provider that sends no usage leaves this at zero rather than
        // guessing — which is why the field is a tally and not a bill.
        let usage = RunUsage::default();

        assert_eq!(usage.total_tokens(), 0);
    }

    #[test]
    fn usage_adds_up_across_runs() {
        // What a shared budget folds: neither run knows about the other.
        let one = RunUsage::default().recording(&usage_report(100, 10));
        let two = RunUsage::default().recording(&usage_report(200, 20));

        assert_eq!(one.plus(two).total_tokens(), 330);
        assert_eq!(one.total_tokens(), 110, "the originals are untouched");
    }
}