lean_ctx/core/
agent_budget.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3
4use serde::{Deserialize, Serialize};
5
6static BUDGETS: Mutex<Option<HashMap<String, AgentBudget>>> = Mutex::new(None);
7static TURN_BUDGETS: Mutex<Option<HashMap<String, TurnBudget>>> = Mutex::new(None);
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct AgentBudget {
11    pub agent_id: String,
12    pub token_limit: usize,
13    pub tokens_consumed: usize,
14    pub reads_count: u32,
15    pub last_reset: String,
16}
17
18#[derive(Debug, Clone)]
19pub struct TurnBudget {
20    pub turn_id: u64,
21    pub tokens_delivered: u64,
22    pub limit: u64,
23    pub started_at: std::time::Instant,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum TurnBudgetStatus {
28    Ok,
29    Warning(u64),
30    Exceeded,
31}
32
33#[derive(Debug, Clone, PartialEq)]
34pub enum BudgetCheckResult {
35    Allowed { remaining: usize },
36    Exceeded { limit: usize, consumed: usize },
37    Warning { remaining: usize, percent_used: f32 },
38}
39
40const WARNING_THRESHOLD: f32 = 0.80;
41
42struct TurnBudgetAssessment {
43    status: TurnBudgetStatus,
44    consumed: u64,
45    projected: u64,
46    limit: u64,
47}
48
49fn with_budgets<F, R>(f: F) -> R
50where
51    F: FnOnce(&mut HashMap<String, AgentBudget>) -> R,
52{
53    let mut guard = BUDGETS
54        .lock()
55        .unwrap_or_else(std::sync::PoisonError::into_inner);
56    let map = guard.get_or_insert_with(HashMap::new);
57    f(map)
58}
59
60fn ensure_entry<'a>(
61    map: &'a mut HashMap<String, AgentBudget>,
62    agent_id: &str,
63) -> &'a mut AgentBudget {
64    map.entry(agent_id.to_string())
65        .or_insert_with(|| AgentBudget {
66            agent_id: agent_id.to_string(),
67            token_limit: usize::MAX,
68            tokens_consumed: 0,
69            reads_count: 0,
70            last_reset: chrono::Utc::now().to_rfc3339(),
71        })
72}
73
74fn assess_turn_budget(budget: &TurnBudget, additional_tokens: u64) -> TurnBudgetAssessment {
75    let projected = budget.tokens_delivered.saturating_add(additional_tokens);
76    let status = if budget.limit > 0 && projected > budget.limit {
77        TurnBudgetStatus::Exceeded
78    } else if budget.limit > 0 && u128::from(projected) * 100 >= u128::from(budget.limit) * 80 {
79        TurnBudgetStatus::Warning(budget.limit.saturating_sub(projected))
80    } else {
81        TurnBudgetStatus::Ok
82    };
83
84    TurnBudgetAssessment {
85        status,
86        consumed: budget.tokens_delivered,
87        projected,
88        limit: budget.limit,
89    }
90}
91
92fn projected_turn_budget_status(
93    agent_id: &str,
94    tokens_to_consume: u64,
95) -> Option<TurnBudgetAssessment> {
96    let mut guard = TURN_BUDGETS
97        .lock()
98        .unwrap_or_else(std::sync::PoisonError::into_inner);
99    guard
100        .get_or_insert_with(HashMap::new)
101        .get(agent_id)
102        .map(|budget| assess_turn_budget(budget, tokens_to_consume))
103}
104
105fn budget_size(value: u64) -> usize {
106    usize::try_from(value).unwrap_or(usize::MAX)
107}
108
109fn turn_budget_check_result(assessment: &TurnBudgetAssessment) -> BudgetCheckResult {
110    match assessment.status {
111        TurnBudgetStatus::Ok => BudgetCheckResult::Allowed {
112            remaining: budget_size(assessment.limit.saturating_sub(assessment.projected)),
113        },
114        TurnBudgetStatus::Warning(remaining) => BudgetCheckResult::Warning {
115            remaining: budget_size(remaining),
116            percent_used: assessment.projected as f32 / assessment.limit as f32,
117        },
118        TurnBudgetStatus::Exceeded => BudgetCheckResult::Exceeded {
119            limit: budget_size(assessment.limit),
120            consumed: budget_size(assessment.consumed),
121        },
122    }
123}
124
125pub fn start_new_turn(agent_id: &str, limit: u64) {
126    let mut guard = TURN_BUDGETS
127        .lock()
128        .unwrap_or_else(std::sync::PoisonError::into_inner);
129    let budgets = guard.get_or_insert_with(HashMap::new);
130    let turn_id = budgets
131        .get(agent_id)
132        .map_or(1, |budget| budget.turn_id.saturating_add(1));
133    budgets.insert(
134        agent_id.to_string(),
135        TurnBudget {
136            turn_id,
137            tokens_delivered: 0,
138            limit,
139            started_at: std::time::Instant::now(),
140        },
141    );
142}
143
144pub fn record_turn_delivery(agent_id: &str, tokens: u64) -> TurnBudgetStatus {
145    let mut guard = TURN_BUDGETS
146        .lock()
147        .unwrap_or_else(std::sync::PoisonError::into_inner);
148    let Some(budget) = guard.get_or_insert_with(HashMap::new).get_mut(agent_id) else {
149        return TurnBudgetStatus::Ok;
150    };
151
152    budget.tokens_delivered = budget.tokens_delivered.saturating_add(tokens);
153    assess_turn_budget(budget, 0).status
154}
155
156pub fn turn_budget_remaining(agent_id: &str) -> Option<u64> {
157    let mut guard = TURN_BUDGETS
158        .lock()
159        .unwrap_or_else(std::sync::PoisonError::into_inner);
160    guard
161        .get_or_insert_with(HashMap::new)
162        .get(agent_id)
163        .map(|budget| budget.limit.saturating_sub(budget.tokens_delivered))
164}
165
166pub fn check_budget(agent_id: &str, tokens_to_consume: usize) -> BudgetCheckResult {
167    let turn_assessment = projected_turn_budget_status(agent_id, tokens_to_consume as u64);
168    if let Some(assessment) = &turn_assessment {
169        if matches!(assessment.status, TurnBudgetStatus::Exceeded) {
170            return turn_budget_check_result(assessment);
171        }
172    }
173
174    with_budgets(|map| {
175        let budget = ensure_entry(map, agent_id);
176        if budget.token_limit == usize::MAX || budget.token_limit == 0 {
177            return turn_assessment.as_ref().map_or(
178                BudgetCheckResult::Allowed {
179                    remaining: usize::MAX,
180                },
181                turn_budget_check_result,
182            );
183        }
184
185        let projected = budget.tokens_consumed.saturating_add(tokens_to_consume);
186        if projected > budget.token_limit {
187            return BudgetCheckResult::Exceeded {
188                limit: budget.token_limit,
189                consumed: budget.tokens_consumed,
190            };
191        }
192
193        let percent_used = projected as f32 / budget.token_limit as f32;
194        let remaining = budget.token_limit.saturating_sub(projected);
195
196        if percent_used >= WARNING_THRESHOLD {
197            BudgetCheckResult::Warning {
198                remaining,
199                percent_used,
200            }
201        } else if let Some(assessment) = &turn_assessment {
202            turn_budget_check_result(assessment)
203        } else {
204            BudgetCheckResult::Allowed { remaining }
205        }
206    })
207}
208
209pub fn record_consumption(agent_id: &str, tokens: usize) {
210    with_budgets(|map| {
211        let budget = ensure_entry(map, agent_id);
212        budget.tokens_consumed = budget.tokens_consumed.saturating_add(tokens);
213        budget.reads_count += 1;
214    });
215    record_turn_delivery(agent_id, tokens as u64);
216}
217
218pub fn get_status(agent_id: &str) -> AgentBudget {
219    with_budgets(|map| ensure_entry(map, agent_id).clone())
220}
221
222pub fn reset(agent_id: &str) {
223    with_budgets(|map| {
224        let budget = ensure_entry(map, agent_id);
225        budget.tokens_consumed = 0;
226        budget.reads_count = 0;
227        budget.last_reset = chrono::Utc::now().to_rfc3339();
228    });
229}
230
231/// Remove an agent's budget entry entirely. Safe only for agents that can no longer
232/// issue reads (finished / dead PID) — a live agent would have its budget silently
233/// reset to 0 on the next check. Bounds the BUDGETS map on long-lived daemons.
234pub fn remove(agent_id: &str) {
235    with_budgets(|map| {
236        map.remove(agent_id);
237    });
238    TURN_BUDGETS
239        .lock()
240        .unwrap_or_else(std::sync::PoisonError::into_inner)
241        .get_or_insert_with(HashMap::new)
242        .remove(agent_id);
243}
244
245pub fn set_limit(agent_id: &str, limit: usize) {
246    with_budgets(|map| {
247        let budget = ensure_entry(map, agent_id);
248        budget.token_limit = if limit == 0 { usize::MAX } else { limit };
249    });
250}
251
252pub fn init_from_config() {
253    let cfg_limit = crate::core::config::Config::load().agent_token_budget;
254    if cfg_limit > 0 {
255        with_budgets(|map| {
256            for budget in map.values_mut() {
257                if budget.token_limit == usize::MAX {
258                    budget.token_limit = cfg_limit;
259                }
260            }
261        });
262    }
263}
264
265pub fn default_limit_from_config() -> usize {
266    let cfg_limit = crate::core::config::Config::load().agent_token_budget;
267    if cfg_limit == 0 {
268        usize::MAX
269    } else {
270        cfg_limit
271    }
272}
273
274#[cfg(test)]
275pub mod tests {
276    use super::*;
277
278    fn test_agent(name: &str) -> String {
279        format!("test_agent_{name}_{:?}", std::thread::current().id())
280    }
281
282    #[test]
283    fn unlimited_budget_always_allows() {
284        let id = test_agent("unlimited");
285        let result = check_budget(&id, 1_000_000);
286        assert!(matches!(result, BudgetCheckResult::Allowed { .. }));
287    }
288
289    #[test]
290    fn set_limit_and_exceed() {
291        let id = test_agent("exceed");
292        set_limit(&id, 1000);
293        record_consumption(&id, 800);
294        let result = check_budget(&id, 300);
295        assert!(matches!(
296            result,
297            BudgetCheckResult::Exceeded {
298                limit: 1000,
299                consumed: 800
300            }
301        ));
302    }
303
304    #[test]
305    fn warning_at_80_percent() {
306        let id = test_agent("warning");
307        set_limit(&id, 1000);
308        record_consumption(&id, 700);
309        let result = check_budget(&id, 100);
310        assert!(matches!(result, BudgetCheckResult::Warning { .. }));
311    }
312
313    #[test]
314    fn reset_clears_consumption() {
315        let id = test_agent("reset");
316        set_limit(&id, 1000);
317        record_consumption(&id, 900);
318        reset(&id);
319        let status = get_status(&id);
320        assert_eq!(status.tokens_consumed, 0);
321        assert_eq!(status.reads_count, 0);
322    }
323
324    #[test]
325    fn zero_limit_means_unlimited() {
326        let id = test_agent("zero");
327        set_limit(&id, 0);
328        let result = check_budget(&id, 1_000_000);
329        assert!(matches!(result, BudgetCheckResult::Allowed { .. }));
330    }
331
332    #[test]
333    fn record_increments_reads_count() {
334        let id = test_agent("reads");
335        record_consumption(&id, 100);
336        record_consumption(&id, 200);
337        let status = get_status(&id);
338        assert_eq!(status.reads_count, 2);
339        assert_eq!(status.tokens_consumed, 300);
340    }
341
342    #[test]
343    fn start_new_turn_resets_delivery_and_increments_id() {
344        let id = test_agent("turn_reset");
345        start_new_turn(&id, 1_000);
346        record_turn_delivery(&id, 600);
347        start_new_turn(&id, 500);
348
349        let budget = TURN_BUDGETS
350            .lock()
351            .unwrap_or_else(std::sync::PoisonError::into_inner)
352            .get_or_insert_with(HashMap::new)
353            .get(&id)
354            .cloned()
355            .unwrap();
356        assert_eq!(budget.turn_id, 2);
357        assert_eq!(budget.tokens_delivered, 0);
358        assert_eq!(budget.limit, 500);
359        assert!(budget.started_at.elapsed().as_secs() < 1);
360        remove(&id);
361    }
362
363    #[test]
364    fn turn_delivery_warns_at_eighty_percent() {
365        let id = test_agent("turn_warning");
366        start_new_turn(&id, 1_000);
367
368        assert_eq!(
369            record_turn_delivery(&id, 800),
370            TurnBudgetStatus::Warning(200)
371        );
372        assert_eq!(turn_budget_remaining(&id), Some(200));
373        remove(&id);
374    }
375
376    #[test]
377    fn turn_delivery_exceeds_limit() {
378        let id = test_agent("turn_exceeded");
379        start_new_turn(&id, 1_000);
380
381        assert_eq!(record_turn_delivery(&id, 1_001), TurnBudgetStatus::Exceeded);
382        assert_eq!(turn_budget_remaining(&id), Some(0));
383        remove(&id);
384    }
385
386    #[test]
387    fn check_budget_uses_active_turn_budget() {
388        let id = test_agent("turn_check");
389        start_new_turn(&id, 1_000);
390        record_turn_delivery(&id, 700);
391
392        assert!(matches!(
393            check_budget(&id, 100),
394            BudgetCheckResult::Warning { remaining: 200, .. }
395        ));
396        assert!(matches!(
397            check_budget(&id, 301),
398            BudgetCheckResult::Exceeded {
399                limit: 1_000,
400                consumed: 700
401            }
402        ));
403        remove(&id);
404    }
405}
lean_ctx/core/agent_budget.rs

lean_ctx/core/
agent_budget.rs
