lean_ctx/core/ocla/
budget.rs

1//! In-memory hierarchical budget enforcement for OCLA request admission.
2
3use std::collections::{HashMap, HashSet};
4
5use chrono::Utc;
6
7use super::types::{OclaError, OclaResult};
8
9/// A budget's organizational level and stable identifier.
10#[derive(Clone, Debug, Eq, Hash, PartialEq)]
11pub enum BudgetScope {
12    Org(String),
13    Team(String),
14    User(String),
15}
16
17/// Daily token and USD caps for one scope.
18#[derive(Clone, Debug, PartialEq)]
19pub struct BudgetLimit {
20    pub scope: BudgetScope,
21    pub max_tokens_per_day: u64,
22    pub max_usd_per_day: f64,
23}
24
25#[derive(Clone, Copy, Debug, Default)]
26struct Consumption {
27    tokens: u64,
28    usd: f64,
29}
30
31/// In-memory daily consumption ledger with explicit org/team/user ancestry.
32#[derive(Clone, Debug, Default)]
33pub struct BudgetLedger {
34    limits: HashMap<BudgetScope, BudgetLimit>,
35    parents: HashMap<BudgetScope, BudgetScope>,
36    consumption: HashMap<(BudgetScope, i64), Consumption>,
37}
38
39impl BudgetLedger {
40    /// Creates an empty budget ledger.
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Adds or replaces the cap for a scope.
46    pub fn set_limit(&mut self, limit: BudgetLimit) {
47        self.limits.insert(limit.scope.clone(), limit);
48    }
49
50    /// Associates a user with a team or a team with an org.
51    pub fn set_parent(&mut self, child: BudgetScope, parent: BudgetScope) {
52        self.parents.insert(child, parent);
53    }
54
55    /// Checks the requested tokens against every configured ancestor cap.
56    pub fn check_budget(&self, scope: &BudgetScope, tokens: u64) -> OclaResult<()> {
57        self.check_budget_with_cost(scope, tokens, 0.0)
58    }
59
60    /// Checks both tokens and USD against every configured ancestor cap.
61    pub fn check_budget_with_cost(
62        &self,
63        scope: &BudgetScope,
64        tokens: u64,
65        usd: f64,
66    ) -> OclaResult<()> {
67        if !usd.is_finite() || usd < 0.0 {
68            return Err(OclaError::InvalidRequest(
69                "budget cost must be finite and non-negative".to_string(),
70            ));
71        }
72
73        let day = current_day();
74        for current in self.lineage(scope)? {
75            let Some(limit) = self.limits.get(&current) else {
76                continue;
77            };
78            let consumed = self
79                .consumption
80                .get(&(current.clone(), day))
81                .copied()
82                .unwrap_or_default();
83            if !limit.max_usd_per_day.is_finite() || limit.max_usd_per_day < 0.0 {
84                return Err(OclaError::InvalidRequest(format!(
85                    "invalid USD budget for {current:?}"
86                )));
87            }
88            if tokens > limit.max_tokens_per_day.saturating_sub(consumed.tokens)
89                || consumed.usd + usd >= limit.max_usd_per_day
90            {
91                return Err(OclaError::InvalidRequest(format!(
92                    "budget exceeded for {current:?}"
93                )));
94            }
95        }
96        Ok(())
97    }
98
99    /// Records usage for the scope and all configured ancestors.
100    pub fn record_consumption(&mut self, scope: &BudgetScope, tokens: u64, usd: f64) {
101        let Ok(lineage) = self.lineage(scope) else {
102            return;
103        };
104        let day = current_day();
105        let usd = if usd.is_finite() && usd > 0.0 {
106            usd
107        } else {
108            0.0
109        };
110        for current in lineage {
111            let consumed = self.consumption.entry((current, day)).or_default();
112            consumed.tokens = consumed.tokens.saturating_add(tokens);
113            consumed.usd += usd;
114        }
115    }
116
117    /// Returns today's consumed tokens for a scope.
118    pub fn consumed_tokens(&self, scope: &BudgetScope) -> u64 {
119        self.consumed(scope).tokens
120    }
121
122    /// Returns today's consumed USD for a scope.
123    pub fn consumed_usd(&self, scope: &BudgetScope) -> f64 {
124        self.consumed(scope).usd
125    }
126
127    fn consumed(&self, scope: &BudgetScope) -> Consumption {
128        self.consumption
129            .get(&(scope.clone(), current_day()))
130            .copied()
131            .unwrap_or_default()
132    }
133
134    fn lineage(&self, scope: &BudgetScope) -> OclaResult<Vec<BudgetScope>> {
135        let mut result = Vec::new();
136        let mut seen = HashSet::new();
137        let mut current = scope.clone();
138        loop {
139            if !seen.insert(current.clone()) {
140                return Err(OclaError::InvalidRequest(
141                    "budget hierarchy contains a cycle".to_string(),
142                ));
143            }
144            result.push(current.clone());
145            let Some(parent) = self.parents.get(&current) else {
146                break;
147            };
148            current = parent.clone();
149        }
150        Ok(result)
151    }
152}
153
154fn current_day() -> i64 {
155    Utc::now().timestamp().div_euclid(86_400)
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    fn limit(scope: BudgetScope, tokens: u64) -> BudgetLimit {
163        BudgetLimit {
164            scope,
165            max_tokens_per_day: tokens,
166            max_usd_per_day: 100.0,
167        }
168    }
169
170    #[test]
171    fn user_budget_rejects_when_team_is_over_limit() {
172        let user = BudgetScope::User("alice".to_string());
173        let team = BudgetScope::Team("platform".to_string());
174        let org = BudgetScope::Org("acme".to_string());
175        let mut ledger = BudgetLedger::new();
176        ledger.set_limit(limit(user.clone(), 1_000));
177        ledger.set_limit(limit(team.clone(), 100));
178        ledger.set_limit(limit(org, 10_000));
179        ledger.set_parent(user.clone(), team.clone());
180        ledger.set_parent(team, BudgetScope::Org("acme".to_string()));
181
182        ledger.record_consumption(&user, 100, 1.0);
183        assert!(matches!(
184            ledger.check_budget(&user, 1),
185            Err(OclaError::InvalidRequest(message)) if message.contains("Team")
186        ));
187    }
188
189    #[test]
190    fn consumption_cascades_to_team_and_org() {
191        let user = BudgetScope::User("alice".to_string());
192        let team = BudgetScope::Team("platform".to_string());
193        let org = BudgetScope::Org("acme".to_string());
194        let mut ledger = BudgetLedger::new();
195        ledger.set_parent(user.clone(), team.clone());
196        ledger.set_parent(team.clone(), org.clone());
197
198        ledger.record_consumption(&user, 42, 2.5);
199        assert_eq!(ledger.consumed_tokens(&user), 42);
200        assert_eq!(ledger.consumed_tokens(&team), 42);
201        assert_eq!(ledger.consumed_tokens(&org), 42);
202        assert!((ledger.consumed_usd(&org) - 2.5).abs() < f64::EPSILON);
203    }
204
205    #[test]
206    fn cost_cap_is_checked_and_invalid_cost_is_rejected() {
207        let scope = BudgetScope::Org("acme".to_string());
208        let mut ledger = BudgetLedger::new();
209        ledger.set_limit(BudgetLimit {
210            scope: scope.clone(),
211            max_tokens_per_day: 100,
212            max_usd_per_day: 5.0,
213        });
214
215        assert!(ledger.check_budget_with_cost(&scope, 1, 6.0).is_err());
216        assert!(ledger.check_budget_with_cost(&scope, 1, -1.0).is_err());
217        ledger.record_consumption(&scope, 1, 5.0);
218        assert!(ledger.check_budget(&scope, 1).is_err());
219    }
220
221    #[test]
222    fn hierarchy_cycles_are_rejected() {
223        let a = BudgetScope::Team("a".to_string());
224        let b = BudgetScope::Org("b".to_string());
225        let mut ledger = BudgetLedger::new();
226        ledger.set_parent(a.clone(), b.clone());
227        ledger.set_parent(b.clone(), a.clone());
228        assert!(ledger.check_budget(&a, 1).is_err());
229    }
230}
lean_ctx/core/ocla/budget.rs

lean_ctx/core/ocla/
budget.rs
