Genarrative/server-rs/crates/module-runtime/src/llm_billing.rs

/// Router quota per mud point, with USD numeric cost multiplied by ten.
pub const LLM_ROUTER_QUOTA_PER_POINT: u64 = 50_000;

#[derive(Debug, PartialEq, Eq)]
pub struct LlmQuotaSettlement {
    pub charged_points: u64,
    pub settled_quota: u64,
}

pub fn calculate_llm_quota_settlement(
    baseline: Option<u64>,
    observed: u64,
    spendable_points: u64,
) -> LlmQuotaSettlement {
    let Some(baseline) = baseline else {
        return LlmQuotaSettlement {
            charged_points: 0,
            settled_quota: observed,
        };
    };
    let points = (observed.saturating_sub(baseline) / LLM_ROUTER_QUOTA_PER_POINT)
        .min(spendable_points)
        .min(i64::MAX as u64);
    LlmQuotaSettlement {
        charged_points: points,
        settled_quota: baseline + points * LLM_ROUTER_QUOTA_PER_POINT,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn initializes_without_charging_history() {
        assert_eq!(
            calculate_llm_quota_settlement(None, 123_456, 100),
            LlmQuotaSettlement {
                charged_points: 0,
                settled_quota: 123_456
            }
        );
    }

    #[test]
    fn carries_fraction_and_replays_without_double_charge() {
        let first = calculate_llm_quota_settlement(Some(0), 61_700, 100);
        assert_eq!(first.charged_points, 1);
        assert_eq!(first.settled_quota, 50_000);
        let next = calculate_llm_quota_settlement(Some(first.settled_quota), 111_700, 99);
        assert_eq!(next.charged_points, 1);
        assert_eq!(next.settled_quota, 100_000);
        assert_eq!(
            calculate_llm_quota_settlement(Some(next.settled_quota), 111_700, 98).charged_points,
            0
        );
        assert_eq!(
            calculate_llm_quota_settlement(Some(next.settled_quota), 61_700, 98).settled_quota,
            100_000
        );
    }

    #[test]
    fn insufficient_balance_only_advances_paid_quota() {
        let first = calculate_llm_quota_settlement(Some(1234), 201_234, 1);
        assert_eq!(first.settled_quota, 51_234);
        assert_eq!(
            calculate_llm_quota_settlement(Some(first.settled_quota), 201_234, 0).settled_quota,
            51_234
        );
        assert_eq!(
            calculate_llm_quota_settlement(Some(first.settled_quota), 201_234, 10).charged_points,
            3
        );
    }

    #[test]
    fn extreme_values_do_not_overflow() {
        let result = calculate_llm_quota_settlement(Some(0), u64::MAX, u64::MAX);
        assert!(result.settled_quota <= u64::MAX - u64::MAX % LLM_ROUTER_QUOTA_PER_POINT);
        assert_eq!(
            calculate_llm_quota_settlement(Some(u64::MAX), 0, u64::MAX).settled_quota,
            u64::MAX
        );
    }
}