disyn_runtime/
budget.rs

1use disyn_core::types::{CostEstimate, ResourceUsage};
2
3pub struct BudgetManager {
4    max_tokens: u64,
5    max_neural_replans: u32,
6    max_repair_attempts: u32,
7    used_tokens: u64,
8    neural_replans: u32,
9    repair_attempts: u32,
10    max_neural_tokens: u64,
11    used_neural_tokens: u64,
12    used_symbolic_tokens: u64,
13}
14
15impl BudgetManager {
16    pub fn new(
17        max_tokens: u64,
18        max_neural_replans: u32,
19        max_repair_attempts: u32,
20        max_neural_tokens: u64,
21    ) -> Self {
22        Self {
23            max_tokens,
24            max_neural_replans,
25            max_repair_attempts,
26            used_tokens: 0,
27            neural_replans: 0,
28            repair_attempts: 0,
29            max_neural_tokens,
30            used_neural_tokens: 0,
31            used_symbolic_tokens: 0,
32        }
33    }
34
35    pub fn can_afford(&self, estimate: &CostEstimate) -> bool {
36        let needed = u64::from(estimate.input_tokens) + u64::from(estimate.output_tokens);
37        self.used_tokens + needed <= self.max_tokens
38    }
39
40    pub fn can_afford_neural(&self, estimate: &CostEstimate) -> bool {
41        let needed = u64::from(estimate.input_tokens) + u64::from(estimate.output_tokens);
42        self.used_neural_tokens + needed <= self.max_neural_tokens
43            && self.used_tokens + needed <= self.max_tokens
44    }
45
46    pub fn record(&mut self, usage: &ResourceUsage) {
47        self.used_tokens += usage.total_tokens;
48        self.used_neural_tokens += usage.neural_tokens;
49        self.used_symbolic_tokens += usage.symbolic_tokens;
50    }
51
52    pub fn can_repair(&self) -> bool {
53        self.repair_attempts < self.max_repair_attempts
54    }
55
56    pub fn record_repair(&mut self) {
57        self.repair_attempts += 1;
58    }
59
60    pub fn can_replan(&self) -> bool {
61        self.neural_replans < self.max_neural_replans
62    }
63
64    pub fn record_replan(&mut self) {
65        self.neural_replans += 1;
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn budget_starts_with_capacity() {
75        let b = BudgetManager::new(1000, 1, 3, 500);
76        assert!(b.can_afford(&CostEstimate {
77            class: None,
78            input_tokens: 500,
79            output_tokens: 200,
80        }));
81    }
82
83    #[test]
84    fn budget_rejects_when_exceeded() {
85        let mut b = BudgetManager::new(100, 1, 3, 100);
86        b.record(&ResourceUsage {
87            total_tokens: 90,
88            symbolic_tokens: 0,
89            neural_tokens: 90,
90            wall_time_ms: 100,
91        });
92        assert!(!b.can_afford(&CostEstimate {
93            class: None,
94            input_tokens: 50,
95            output_tokens: 50,
96        }));
97    }
98
99    #[test]
100    fn replan_attempts_decrement() {
101        let mut b = BudgetManager::new(1000, 2, 3, 500);
102        assert!(b.can_replan());
103        b.record_replan();
104        b.record_replan();
105        assert!(!b.can_replan());
106    }
107
108    #[test]
109    fn repair_attempts_decrement() {
110        let mut b = BudgetManager::new(1000, 1, 3, 500);
111        assert!(b.can_repair());
112        b.record_repair();
113        b.record_repair();
114        b.record_repair();
115        assert!(!b.can_repair());
116    }
117
118    #[test]
119    fn neural_budget_rejects_when_class_exceeded() {
120        use disyn_core::types::CostClass;
121        let mut b = BudgetManager::new(10_000, 1, 3, 100);
122        b.record(&ResourceUsage {
123            total_tokens: 90,
124            symbolic_tokens: 0,
125            neural_tokens: 90,
126            wall_time_ms: 10,
127        });
128        assert!(!b.can_afford_neural(&CostEstimate {
129            class: Some(CostClass::Neural),
130            input_tokens: 50,
131            output_tokens: 50,
132        }));
133    }
134}
disyn_runtime/budget.rs

disyn_runtime/
budget.rs
