Skip to main content

aidens_budget_kit/
lib.rs

1//! Budget, retry, fanout, and stop-rule guards.
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct BudgetV1 {
5    pub max_tool_calls: u32,
6    pub max_retries: u32,
7    pub max_turn_millis: u64,
8}
9
10impl Default for BudgetV1 {
11    fn default() -> Self {
12        Self {
13            max_tool_calls: 8,
14            max_retries: 2,
15            max_turn_millis: 30_000,
16        }
17    }
18}
19
20impl BudgetV1 {
21    pub fn allows_tool_call(&self, tool_calls_so_far: u32, requested_tool_calls: u32) -> bool {
22        tool_calls_so_far.saturating_add(requested_tool_calls) <= self.max_tool_calls
23    }
24
25    pub fn allows_retry(&self, retries_so_far: u32) -> bool {
26        retries_so_far < self.max_retries
27    }
28
29    pub fn deadline_exceeded(&self, elapsed_millis: u64) -> bool {
30        elapsed_millis > self.max_turn_millis
31    }
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    #[test]
39    fn tool_call_budget_is_cumulative() {
40        let budget = BudgetV1 {
41            max_tool_calls: 2,
42            max_retries: 1,
43            max_turn_millis: 1000,
44        };
45
46        assert!(budget.allows_tool_call(1, 1));
47        assert!(!budget.allows_tool_call(2, 1));
48    }
49
50    #[test]
51    fn deadline_check_is_explicit() {
52        let budget = BudgetV1 {
53            max_tool_calls: 1,
54            max_retries: 1,
55            max_turn_millis: 10,
56        };
57
58        assert!(!budget.deadline_exceeded(10));
59        assert!(budget.deadline_exceeded(11));
60    }
61}