Skip to main content

mcp_usage_core/
limits.rs

1//! Pure quota and spend-cap admission decisions.
2//!
3//! The edge or control plane owns the authoritative counters. This module only
4//! performs the overflow-safe decision so every deployment applies identical
5//! boundary semantics.
6
7/// Tenant limits. `None` means unbounded.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
9pub struct Limits {
10    /// Maximum metered units in the control plane's current window.
11    pub max_units: Option<u64>,
12    /// Maximum spend in millionths of the configured currency unit.
13    pub max_spend_micros: Option<u64>,
14}
15
16/// Usage already committed in the current window.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
18pub struct Usage {
19    /// Metered units already committed.
20    pub units: u64,
21    /// Spend already committed, in millionths of a currency unit.
22    pub spend_micros: u64,
23}
24
25/// Successful reservation or the reason admission was denied.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum LimitDecision {
28    /// Admission succeeds with these new committed totals.
29    Allowed(Usage),
30    /// Admission would violate a configured limit.
31    Rejected(LimitReason),
32}
33
34impl LimitDecision {
35    /// Whether the reservation may proceed.
36    #[must_use]
37    pub const fn is_allowed(self) -> bool {
38        matches!(self, Self::Allowed(_))
39    }
40}
41
42/// Why a usage reservation was rejected.
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum LimitReason {
45    /// The unit quota would be exceeded.
46    QuotaExceeded,
47    /// The monetary spend cap would be exceeded.
48    SpendCapExceeded,
49    /// Calculating the requested total exceeded the representable range.
50    ArithmeticOverflow,
51}
52
53/// Atomically assess a proposed usage increment against quota and spend caps.
54///
55/// `unit_price_micros` is the monetary price of one metered unit. Exact boundary
56/// values are allowed: a limit rejects only when the new total is greater.
57#[must_use]
58pub fn assess_limits(
59    current: Usage,
60    requested_units: u64,
61    unit_price_micros: u64,
62    limits: Limits,
63) -> LimitDecision {
64    let Some(units) = current.units.checked_add(requested_units) else {
65        return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
66    };
67    if limits.max_units.is_some_and(|limit| units > limit) {
68        return LimitDecision::Rejected(LimitReason::QuotaExceeded);
69    }
70
71    let Some(incremental_spend) = requested_units.checked_mul(unit_price_micros) else {
72        return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
73    };
74    let Some(spend_micros) = current.spend_micros.checked_add(incremental_spend) else {
75        return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
76    };
77    if limits
78        .max_spend_micros
79        .is_some_and(|limit| spend_micros > limit)
80    {
81        return LimitDecision::Rejected(LimitReason::SpendCapExceeded);
82    }
83
84    LimitDecision::Allowed(Usage {
85        units,
86        spend_micros,
87    })
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn exact_boundaries_are_allowed() {
96        let decision = assess_limits(
97            Usage {
98                units: 8,
99                spend_micros: 800,
100            },
101            2,
102            100,
103            Limits {
104                max_units: Some(10),
105                max_spend_micros: Some(1_000),
106            },
107        );
108        assert_eq!(
109            decision,
110            LimitDecision::Allowed(Usage {
111                units: 10,
112                spend_micros: 1_000
113            })
114        );
115    }
116
117    #[test]
118    fn quota_is_checked_before_spend() {
119        let decision = assess_limits(
120            Usage::default(),
121            11,
122            1_000,
123            Limits {
124                max_units: Some(10),
125                max_spend_micros: Some(1),
126            },
127        );
128        assert_eq!(
129            decision,
130            LimitDecision::Rejected(LimitReason::QuotaExceeded)
131        );
132    }
133
134    #[test]
135    fn spend_cap_rejects_the_first_unit_over() {
136        let decision = assess_limits(
137            Usage {
138                units: 2,
139                spend_micros: 200,
140            },
141            1,
142            101,
143            Limits {
144                max_units: None,
145                max_spend_micros: Some(300),
146            },
147        );
148        assert_eq!(
149            decision,
150            LimitDecision::Rejected(LimitReason::SpendCapExceeded)
151        );
152    }
153
154    #[test]
155    fn arithmetic_overflow_never_wraps_into_an_allowance() {
156        assert_eq!(
157            assess_limits(Usage::default(), u64::MAX, 2, Limits::default()),
158            LimitDecision::Rejected(LimitReason::ArithmeticOverflow)
159        );
160        assert_eq!(
161            assess_limits(
162                Usage {
163                    units: u64::MAX,
164                    spend_micros: 0
165                },
166                1,
167                0,
168                Limits::default()
169            ),
170            LimitDecision::Rejected(LimitReason::ArithmeticOverflow)
171        );
172    }
173}